Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

First pass at adding support for scaling in rand/randn. #633

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

First pass at adding support for scaling in rand/randn. #633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
33ce120
First pass at adding support for scaling in rand/randn.
a3ahmad Dec 8, 2023
5f6dce1
Remove f prefix for strings without arguments.
a3ahmad Dec 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion theseus/geometry/lie_group.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
# LICENSE file in the root directory of this source tree.

import abc
from typing import Any, List, Optional, Tuple, cast
from typing import Any, List, Optional, Tuple, Union, cast

import torch

Expand Down Expand Up @@ -59,6 +59,7 @@ def dof(self) -> int:
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
Expand All @@ -70,6 +71,7 @@ def rand(
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
Expand Down
28 changes: 24 additions & 4 deletions theseus/geometry/point_types.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,12 +42,17 @@ def __init__(
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
) -> "Point2":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (2,):
raise ValueError("The scale must be None, a float, or a 2 element 1D tensor.")
elif scale is None:
scale = 1.0
return Point2(
tensor=torch.rand(
size[0],
Expand All @@ -56,19 +61,24 @@ def rand(
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
)

@staticmethod
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
) -> "Point2":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (2,):
raise ValueError("The scale must be None, a float, or a 2 element 1D tensor.")
elif scale is None:
scale = 1.0
return Point2(
tensor=torch.randn(
size[0],
Expand All @@ -77,7 +87,7 @@ def randn(
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
)

@staticmethod
Expand Down Expand Up @@ -144,12 +154,17 @@ def __init__(
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
) -> "Point3":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (3,):
raise ValueError("The scale must be None, a float, or a 3D vector.")
elif scale is None:
scale = 1.0
return Point3(
tensor=torch.rand(
size[0],
Expand All @@ -158,19 +173,24 @@ def rand(
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
)

@staticmethod
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
) -> "Point3":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (3,):
raise ValueError("The scale must be None, a float, or a 3D vector.")
elif scale is None:
scale = 1.0
return Point3(
tensor=torch.randn(
size[0],
Expand All @@ -179,7 +199,7 @@ def randn(
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
)

@staticmethod
Expand Down
14 changes: 12 additions & 2 deletions theseus/geometry/se2.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,20 +46,25 @@ def __init__(
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SE2":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (3,):
raise ValueError("The scale must be None, a float, or a 3 element 1D tensor.")
elif scale is None:
scale = 1.0
x_y_theta = torch.rand(
size[0],
3,
generator=generator,
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
x_y_theta[:, 2] = 2 * theseus.constants.PI * (x_y_theta[:, 2] - 0.5)

return SE2(x_y_theta=x_y_theta)
Expand All @@ -68,20 +73,25 @@ def rand(
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SE2":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (3,):
raise ValueError("The scale must be None, a float, or a 3 element 1D tensor.")
elif scale is None:
scale = 1.0
x_y_theta = torch.randn(
size[0],
3,
generator=generator,
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
x_y_theta[:, 2] *= theseus.constants.PI

return SE2(x_y_theta=x_y_theta)
Expand Down
12 changes: 12 additions & 0 deletions theseus/geometry/se3.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,15 +45,21 @@ def __init__(
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SE3":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (4,):
raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.")
elif scale is None:
scale = 1.0
tensor = SE3_base.rand(
*size,
generator=generator,
scale=scale,
dtype=dtype,
device=device,
requires_grad=requires_grad,
Expand All @@ -64,15 +70,21 @@ def rand(
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SE3":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (4,):
raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.")
elif scale is None:
scale = 1.0
tensor = SE3_base.randn(
*size,
generator=generator,
scale=scale,
dtype=dtype,
device=device,
requires_grad=requires_grad,
Expand Down
14 changes: 12 additions & 2 deletions theseus/geometry/so2.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,12 +50,17 @@ def __init__(
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SO2":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (1,):
raise ValueError("The scale must be None, or a float, or a single element tensor.")
elif scale is None:
scale = 1.0
return SO2.exp_map(
2
* theseus.constants.PI
Expand All @@ -66,20 +71,25 @@ def rand(
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
- theseus.constants.PI
)

@staticmethod
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SO2":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (1,):
raise ValueError("The scale must be None, or a float, or a single element tensor.")
elif scale is None:
scale = 1.0
return SO2.exp_map(
theseus.constants.PI
* torch.randn(
Expand All @@ -89,7 +99,7 @@ def randn(
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
)

@staticmethod
Expand Down
12 changes: 12 additions & 0 deletions theseus/geometry/so3.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,15 +45,21 @@ def __init__(
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SO3":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (4,):
raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.")
elif scale is None:
scale = 1.0
tensor = SO3_base.rand(
*size,
generator=generator,
scale=scale,
dtype=dtype,
device=device,
requires_grad=requires_grad,
Expand All @@ -64,15 +70,21 @@ def rand(
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: theseus.constants.DeviceType = None,
requires_grad: bool = False,
) -> "SO3":
if len(size) != 1:
raise ValueError("The size should be 1D.")
if isinstance(scale, torch.Tensor) and scale.shape != (4,):
raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.")
elif scale is None:
scale = 1.0
tensor = SO3_base.randn(
*size,
generator=generator,
scale=scale,
dtype=dtype,
device=device,
requires_grad=requires_grad,
Expand Down
14 changes: 12 additions & 2 deletions theseus/geometry/vector.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,40 +46,50 @@ def dof(self) -> int:
def rand(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
) -> "Vector":
if len(size) != 2:
raise ValueError("The size should be 2D.")
if isinstance(scale, torch.Tensor) and (scale.shape[0] != size[0] or scale.shape[1] != size[1]):
raise ValueError(f"The scale must be None, or a float, or a tensor of shape {size}.")
elif scale is None:
scale = 1.0
return Vector(
tensor=torch.rand(
size,
generator=generator,
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
)

@staticmethod
def randn(
*size: int,
generator: Optional[torch.Generator] = None,
scale: Optional[Union[float, torch.Tensor]] = None,
dtype: Optional[torch.dtype] = None,
device: DeviceType = None,
requires_grad: bool = False,
) -> "Vector":
if len(size) != 2:
raise ValueError("The size should be 2D.")
if isinstance(scale, torch.Tensor) and (scale.shape[0] != size[0] or scale.shape[1] != size[1]):
raise ValueError(f"The scale must be None, or a float, or a tensor of shape {size}.")
elif scale is None:
scale = 1.0
return Vector(
tensor=torch.randn(
size,
generator=generator,
dtype=dtype,
device=device,
requires_grad=requires_grad,
)
) * scale
)

def __repr__(self) -> str:
Expand Down
Loading