Skip to content

added cuda support for all pytorch mlg operations #25

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

Merged
angadbajwa merged 6 commits into from
May 1, 2026
Merged

added cuda support for all pytorch mlg operations #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
9d55745
added cuda support for all pytorch mlg operations. changed testing st…
angadbajwa Sep 2, 2025
180360d
perform cuda testing pending GPU availability
angadbajwa Apr 30, 2026
157b0b4
Merge branch 'main' into device_complex_hotfix
angadbajwa Apr 30, 2026
cfdc173
check device count as tests appear to be returning CUDA availability …
angadbajwa Apr 30, 2026
95c57a2
marked tests to be skipped if cuda not available
angadbajwa Apr 30, 2026
8cc3b9d
fixed SE2 device allocations and some base class cleanup
angadbajwa May 1, 2026
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: 2 additions & 2 deletions pymlg/torch/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ def adjoint(Xi):
raise NotImplementedError()

@classmethod
def identity(cls):
def identity(cls, device='cpu'):
"""
Returns an identity matrix of the group.

Expand All @@ -138,7 +138,7 @@ def identity(cls):
torch.Tensor
Identity matrix of the group with shape `(N, n, n)`.
"""
return torch.eye(cls.matrix_size, cls.matrix_size).unsqueeze(0)
return torch.eye(cls.matrix_size, cls.matrix_size, device=device).unsqueeze(0)

@staticmethod
def left_jacobian(x):
Expand Down
33 changes: 18 additions & 15 deletions pymlg/torch/se2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,21 @@ class SE2(MatrixLieGroupTorch):
matrix_size = 3

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
phi = torch.rand(N, 1, 1) * 2 * torch.pi
r = torch.randn(N, 2, 1)
C = SO2.Exp(phi)
return SE2.from_components(C, r)
return SE2.from_components(C, r).to(device)

@staticmethod
def from_components(C, r):
"""
Construct an SE(2) matrix from a rotation matrix and translation vector.
"""
T = torch.zeros(C.shape[0], 3, 3, dtype=C.dtype)

# first, confirm that both components are allocated on the same device
assert C.device == r.device, "Components must be on the same device for SE2.from_components"
Copy link

Copilot AI Apr 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using assert for input validation is unsafe (asserts can be stripped with python -O). Prefer raising a ValueError/RuntimeError for device mismatch so this remains enforced in production.

Suggested change
assert C.device == r.device, "Components must be on the same device for SE2.from_components"
if C.device != r.device:
raise ValueError("Components must be on the same device for SE2.from_components")

Copilot uses AI. Check for mistakes.
T = torch.zeros(C.shape[0], 3, 3, dtype=C.dtype, device=C.device)

T[:, 0:2, 0:2] = C
T[:, 0:2, 2] = r.view(-1, 2)
Expand All @@ -45,7 +48,7 @@ def wedge(xi):
phi = xi[:, 0]
xi_r = xi[:, 1:]
Xi_phi = SO2.wedge(phi)
Xi = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype)
Xi = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype, device=xi.device)
Xi[:, 0:2, 0:2] = Xi_phi
Xi[:, 0:2, 2] = xi_r.view(-1, 2)
return Xi
Expand All @@ -72,17 +75,17 @@ def log(T):
Xi_phi = SO2.log(T[:, 0:2, 0:2])
r = T[:, 0:2, 2].unsqueeze(2)
xi_r = SE2.V_matrix_inv(SO2.vee(Xi_phi)) @ r
Xi = torch.zeros(T.shape[0], 3, 3, dtype=T.dtype)
Xi = torch.zeros(T.shape[0], 3, 3, dtype=T.dtype, device=T.device)
Xi[:, 0:2, 0:2] = Xi_phi
Xi[:, 0:2, 2] = xi_r.squeeze(2)
return Xi

@staticmethod
def odot(b):

X = torch.zeros(b.shape[0], 3, 3, dtype=b.dtype)
X = torch.zeros(b.shape[0], 3, 3, dtype=b.dtype, device=b.device)
X[:, 0:2, 0] = SO2.odot(b[:, :2]).squeeze(2)
X[:, 0:2, 1:3] = batch_eye(b.shape[0], 2, 2) * b[:, 2].unsqueeze(2)
X[:, 0:2, 1:3] = batch_eye(b.shape[0], 2, 2, device=b.device) * b[:, 2].unsqueeze(2)

return X

Expand All @@ -103,7 +106,7 @@ def left_jacobian(xi):
large_angle_mask = small_angle_mask.logical_not()
large_angle_inds = large_angle_mask.nonzero(as_tuple=True)[0]

J = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype)
J = torch.zeros(xi.shape[0], 3, 3, dtype=xi.dtype, device=xi.device)

if small_angle_inds.numel():
A = (1 - 1.0 / 6.0 * phi_sq[small_angle_inds]).view(-1)
Expand Down Expand Up @@ -144,9 +147,9 @@ def adjoint(T):
r = T[:, 0:2, 2]

# build Om matrix manually (will this break the DAG?)
Om = torch.Tensor([[0, -1], [1, 0]]).repeat(T.shape[0], 1, 1)
Om = torch.tensor([[0, -1], [1, 0]], dtype=T.dtype, device=T.device).repeat(T.shape[0], 1, 1)

A = torch.zeros(T.shape[0], 3, 3, dtype=T.dtype)
A = torch.zeros(T.shape[0], 3, 3, dtype=T.dtype, device=T.device)
A[:, 0, 0] = 1
A[:, 1:, 0] = -(Om @ r.unsqueeze(2)).squeeze(2)
A[:, 1:, 1:] = C
Expand All @@ -155,7 +158,7 @@ def adjoint(T):

@staticmethod
def adjoint_algebra(Xi):
A = torch.zeros(Xi.shape[0], 3, 3, dtype=Xi.dtype)
A = torch.zeros(Xi.shape[0], 3, 3, dtype=Xi.dtype, device=Xi.device)
A[:, 1, 0] = Xi[:, 1, 2]
A[:, 2, 0] = -Xi[:, 0, 2]
A[:, 1:, 1:] = Xi[:, 0:2, 0:2]
Expand All @@ -171,7 +174,7 @@ def V_matrix(phi):
large_angle_mask = small_angle_mask.logical_not()
large_angle_inds = large_angle_mask.nonzero(as_tuple=True)[0]

V = batch_eye(phi.shape[0], 2, 2, dtype=phi.dtype)
V = batch_eye(phi.shape[0], 2, 2, device=phi.device, dtype=phi.dtype)

if small_angle_inds.numel():
V[small_angle_inds] += .5 * SO2.wedge(phi[small_angle_inds])
Expand All @@ -180,7 +183,7 @@ def V_matrix(phi):
s = torch.sin(phi[large_angle_inds])
c = torch.cos(phi[large_angle_inds])

V[large_angle_inds] = V[large_angle_inds] * (s / phi[large_angle_inds]).view(-1, 1, 1) + ((1 - c) / phi[large_angle_inds]).view(-1, 1, 1) * SO2.wedge(torch.ones(phi[large_angle_inds].shape[0]))
V[large_angle_inds] = V[large_angle_inds] * (s / phi[large_angle_inds]).view(-1, 1, 1) + ((1 - c) / phi[large_angle_inds]).view(-1, 1, 1) * SO2.wedge(torch.ones(phi[large_angle_inds].shape[0], device=phi.device, dtype=phi.dtype))

return V

Expand All @@ -193,15 +196,15 @@ def V_matrix_inv(phi):
large_angle_mask = small_angle_mask.logical_not()
large_angle_inds = large_angle_mask.nonzero(as_tuple=True)[0]

V_inv = batch_eye(phi.shape[0], 2, 2, dtype=phi.dtype)
V_inv = batch_eye(phi.shape[0], 2, 2, device=phi.device, dtype=phi.dtype)

if small_angle_inds.numel():
V_inv[small_angle_inds] -= .5 * SO2.wedge(phi[small_angle_inds])

if large_angle_inds.numel():
half_angle = phi[large_angle_inds] / 2
cot_half_angle = 1 / torch.tan(half_angle)
V_inv[large_angle_inds] = V_inv[large_angle_inds] * half_angle * cot_half_angle - half_angle * SO2.wedge(torch.ones(half_angle.shape[0]))
V_inv[large_angle_inds] = V_inv[large_angle_inds] * half_angle * cot_half_angle - half_angle * SO2.wedge(torch.ones(half_angle.shape[0], device=phi.device, dtype=phi.dtype))

return V_inv

Expand Down
29 changes: 16 additions & 13 deletions pymlg/torch/se23.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class SE23(MatrixLieGroupTorch):
matrix_size = 5

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
"""
Generates a random batch of SE_2(3) matricies.

Expand All @@ -39,7 +39,7 @@ def random(N=1):

C = SO3.Exp(phi)

return SE23.from_components(C, v, r)
return SE23.from_components(C, v, r).to(device)

@staticmethod
def from_components(C: torch.Tensor, v: torch.Tensor, r: torch.Tensor):
Expand All @@ -65,8 +65,11 @@ def from_components(C: torch.Tensor, v: torch.Tensor, r: torch.Tensor):
# firstly, check that batch dimension for all 3 components matches
if not (C.shape[0] == v.shape[0] == r.shape[0]):
raise ValueError("Batch dimension for SE_2(3) components don't match.")

# check that all components are on the same device
assert C.device == v.device == r.device, "Components must be on the same device for SE23.from_components."
Copy link

Copilot AI Apr 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using assert for input validation is unsafe here (asserts can be stripped with python -O). Prefer raising a ValueError/RuntimeError for device mismatch so this remains enforced in optimized runs.

Suggested change
assert C.device == v.device == r.device, "Components must be on the same device for SE23.from_components."
if not (C.device == v.device == r.device):
raise ValueError("Components must be on the same device for SE23.from_components.")

Copilot uses AI. Check for mistakes.

X = batch_eye(C.shape[0], 5, 5, dtype=C.dtype)
X = batch_eye(C.shape[0], 5, 5, device=C.device, dtype=C.dtype)

X[:, 0:3, 0:3] = C
X[:, 0:3, 3] = v.squeeze(2)
Expand Down Expand Up @@ -112,7 +115,7 @@ def wedge(xi: torch.Tensor):
Xi = torch.cat(
(SO3.cross(xi_phi), xi_v, xi_r), dim=2
) # this yields a (N, 3, 5) matrix that must now be blocked with a (2, 5) batched matrix
block = torch.zeros(xi_phi.shape[0], 2, 5)
block = torch.zeros(xi_phi.shape[0], 2, 5, device=xi.device, dtype=xi.dtype)
return torch.cat((Xi, block), dim=1)

@staticmethod
Expand Down Expand Up @@ -152,13 +155,13 @@ def log(X):
Xi = torch.cat(
(SO3.cross(phi), v, r), dim=2
) # this yields a (N, 3, 5) matrix that must now be blocked with a (2, 5) batched matrix
block = torch.zeros(X.shape[0], 2, 5)
block = torch.zeros(X.shape[0], 2, 5, device=X.device, dtype=X.dtype)
return torch.cat((Xi, block), dim=1)

@staticmethod
def adjoint(X):
C, v, r = SE23.to_components(X)
O = torch.zeros(v.shape[0], 3, 3)
O = torch.zeros(v.shape[0], 3, 3, device=X.device, dtype=X.dtype)

# creating block matrix
b1 = torch.cat((C, O, O), dim=2)
Expand All @@ -168,7 +171,7 @@ def adjoint(X):

@staticmethod
def adjoint_algebra(Xi):
A = torch.zeros(Xi.shape[0], 9, 9)
A = torch.zeros(Xi.shape[0], 9, 9, dtype=Xi.dtype, device=Xi.device)
A[:, 0:3, 0:3] = Xi[:, 0:3, 0:3]
A[:, 3:6, 0:3] = SO3.wedge(Xi[:, 0:3, 3])
A[:, 3:6, 3:6] = Xi[:, 0:3, 0:3]
Expand All @@ -177,14 +180,14 @@ def adjoint_algebra(Xi):
return A

@staticmethod
def identity(N=1):
return batch_eye(N, 5, 5)
def identity(device, N=1, dtype=torch.float64):
Copy link

Copilot AI Apr 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SE23.identity now requires a "device" argument, which breaks the MatrixLieGroupTorch.identity() interface and existing call sites that expect identity() with no args. Consider making device optional (defaulting to 'cpu' or inferred) to preserve compatibility.

Suggested change
def identity(device, N=1, dtype=torch.float64):
def identity(device='cpu', N=1, dtype=torch.float64):

Copilot uses AI. Check for mistakes.
return batch_eye(N, 5, 5, device=device, dtype=dtype)

@staticmethod
def odot(xi : torch.Tensor):
X = torch.zeros(xi.shape[0], 5, 9)
X = torch.zeros(xi.shape[0], 5, 9, dtype=xi.dtype, device=xi.device)
X[:, 0:4, 0:6] = SE3.odot(xi[:, 0:4])
X[:, 0:3, 6:9] = xi[:, 4] * batch_eye(xi.shape[0], 3, 3)
X[:, 0:3, 6:9] = xi[:, 4] * batch_eye(xi.shape[0], 3, 3, device=xi.device, dtype=xi.dtype)
return X

@staticmethod
Expand All @@ -193,7 +196,7 @@ def left_jacobian(xi):
xi_v = xi[:, 3:6]
xi_r = xi[:, 6:9]

J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype, device=xi.device)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE23._small_angle_tol
Expand Down Expand Up @@ -242,7 +245,7 @@ def left_jacobian_inv(xi):
xi_v = xi[:, 3:6]
xi_r = xi[:, 6:9]

J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 9, 9, dtype=xi.dtype, device=xi.device)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE23._small_angle_tol
Expand Down
29 changes: 16 additions & 13 deletions pymlg/torch/se3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def _left_jacobian_Q_matrix(xi_phi, xi_rho):
return Q.squeeze_()

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
"""
Generates a random batch of SE_(3) matricies.

Expand All @@ -70,7 +70,7 @@ def random(N=1):

C = SO3.Exp(phi)

return SE3.from_components(C, r)
return SE3.from_components(C, r).to(device)

@staticmethod
def from_components(C: torch.Tensor, r: torch.Tensor):
Expand All @@ -94,8 +94,11 @@ def from_components(C: torch.Tensor, r: torch.Tensor):
# firstly, check that batch dimension for all 3 components matches
if not (C.shape[0] == r.shape[0]):
raise ValueError("Batch dimension for SE(3) components don't match.")

# then, check that both components are on the same device
assert C.device == r.device, "Components must be on the same device for SE3.from_components."
Copy link

Copilot AI Apr 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using assert for input validation is unsafe here (asserts can be stripped with python -O). Since you already raise ValueError for shape mismatches, use a ValueError/RuntimeError for device mismatches as well to ensure this check always runs in production.

Suggested change
assert C.device == r.device, "Components must be on the same device for SE3.from_components."
if C.device != r.device:
raise ValueError("Components must be on the same device for SE3.from_components.")

Copilot uses AI. Check for mistakes.

X = batch_eye(C.shape[0], 4, 4, dtype = C.dtype)
X = batch_eye(C.shape[0], 4, 4, device=C.device, dtype = C.dtype)

X[:, 0:3, 0:3] = C
X[:, 0:3, 3] = r.squeeze(2)
Expand Down Expand Up @@ -140,7 +143,7 @@ def wedge(xi: torch.Tensor):
) # this yields a (N, 3, 4) matrix that must now be blocked with a (1, 4) batched matrix

# generating a (N, 1, 4) batched matrix to append
b1 = torch.tensor([0, 0, 0, 0], dtype = xi.dtype).reshape(1, 1, 4)
b1 = torch.tensor([0, 0, 0, 0], dtype = xi.dtype, device=xi.device).reshape(1, 1, 4)
block = b1.repeat(Xi.shape[0], 1, 1)

return torch.cat((Xi, block), dim=1)
Expand Down Expand Up @@ -180,22 +183,22 @@ def log(X : torch.Tensor):
) # this yields a (N, 3, 4) matrix that must now be blocked with a (1, 4) batched matrix

# generating a (N, 1, 4) batched matrix to append
b1 = torch.tensor([0, 0, 0, 0]).reshape(1, 1, 4)
b1 = torch.tensor([0, 0, 0, 0], device=X.device).reshape(1, 1, 4)
Copy link

Copilot AI Apr 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

b1 is created without specifying dtype, so it will default to an integer/float32 tensor and can cause dtype mismatch when concatenating with Xi (often float64). On CUDA this can raise an error. Construct b1 with the same dtype/device as X (or Xi) to keep concatenation valid.

Suggested change
b1 = torch.tensor([0, 0, 0, 0], device=X.device).reshape(1, 1, 4)
b1 = torch.tensor([0, 0, 0, 0], dtype=Xi.dtype, device=Xi.device).reshape(1, 1, 4)

Copilot uses AI. Check for mistakes.
block = b1.repeat(Xi.shape[0], 1, 1)

return torch.cat((Xi, block), dim=1)

@staticmethod
def odot(b : torch.Tensor):
X = torch.zeros(b.shape[0], 4, 6, dtype=b.dtype)
X = torch.zeros(b.shape[0], 4, 6, dtype=b.dtype, device=b.device)
X[:, 0:3, 0:3] = SO3.odot(b[0:3])
X[:, 0:3, 3:6] = b[:, 3] * batch_eye(b.shape[0], 3, 3, dtype=b.dtype)
X[:, 0:3, 3:6] = b[:, 3] * batch_eye(b.shape[0], 3, 3, dtype=b.dtype, device=b.device)
Comment on lines +193 to +195
Copy link

Copilot AI Apr 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In SE3.odot, b[0:3] slices the batch dimension, not the first three vector components. This gives incorrect results when batch size > 1 (and also passes a (N,4,1) tensor into SO3.odot). Use component slicing on the second dimension (e.g., b[:, 0:3]) so the odot operator is computed from the rotational part of each batch element.

Copilot uses AI. Check for mistakes.
return X

@staticmethod
def adjoint(X):
C, r = SE3.to_components(X)
O = torch.zeros(r.shape[0], 3, 3)
O = torch.zeros(r.shape[0], 3, 3, device=X.device)

# creating block matrix
b1 = torch.cat((C, O), dim=2)
Expand All @@ -204,22 +207,22 @@ def adjoint(X):

@staticmethod
def adjoint_algebra(Xi):
A = torch.zeros(Xi.shape[0], 6, 6, dtype=Xi.dtype)
A = torch.zeros(Xi.shape[0], 6, 6, dtype=Xi.dtype, device=Xi.device)
A[:, 0:3, 0:3] = Xi[:, 0:3, 0:3]
A[:, 3:6, 0:3] = SO3.wedge(Xi[:, 0:3, 3])
A[:, 3:6, 3:6] = Xi[:, 0:3, 0:3]
return A

@staticmethod
def identity(N=1, dtype=torch.float32):
return batch_eye(N, 4, 4, dtype=dtype)
def identity(device, N=1, dtype=torch.float64):
Copy link

Copilot AI Apr 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SE3.identity now requires a "device" argument, which breaks the MatrixLieGroupTorch.identity() interface and existing call sites that expect identity() to be callable with no args. Consider making device optional (defaulting to 'cpu' or inferred) to preserve compatibility.

Suggested change
def identity(device, N=1, dtype=torch.float64):
def identity(device=None, N=1, dtype=torch.float64):

Copilot uses AI. Check for mistakes.
return batch_eye(N, 4, 4, device=device, dtype=dtype)

@staticmethod
def left_jacobian(xi):
xi_phi = xi[:, 0:3]
xi_r = xi[:, 3:6]

J_left = batch_eye(xi.shape[0], 6, 6, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 6, 6, device=xi.device, dtype=xi.dtype)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE3._small_angle_tol
Expand Down Expand Up @@ -256,7 +259,7 @@ def left_jacobian_inv(xi):
xi_phi = xi[:, 0:3]
xi_r = xi[:, 3:6]

J_left = batch_eye(xi.shape[0], 6, 6, dtype=xi.dtype)
J_left = batch_eye(xi.shape[0], 6, 6, device=xi.device, dtype=xi.dtype)

small_angle_mask = is_close(
torch.linalg.norm(xi_phi, dim=1), 0.0, SE3._small_angle_tol
Expand Down
4 changes: 2 additions & 2 deletions pymlg/torch/so2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ class SO2(MatrixLieGroupTorch):
matrix_size = 2

@staticmethod
def random(N=1):
def random(N=1, device='cpu'):
"""
Generates a random batch of SO_(2) matricies.

Expand All @@ -21,7 +21,7 @@ def random(N=1):
batch size, by default 1
"""
phi = torch.rand(N, 1) * 2 * torch.pi
return SO2.Exp(phi)
return SO2.Exp(phi).to(device)

@staticmethod
def wedge(phi):
Expand Down
Loading
Loading