Merge pull request pypose#284 from pypose/lie_op_vmap

Allow some LieTensor operation workable under `func.vmap`

pypose/basics/ops.py

@@ -1,7 +1,7 @@
 import math, torch
 
 
-def pm(input, *, out=None):
+def pm(input):
     r'''
     Returns plus or minus (:math:`\pm`) states for tensor.
 
@@ -21,9 +21,7 @@ def pm(input, *, out=None):
         >>> pp.pm(torch.tensor([0.1, 0, -0.2], dtype=torch.float64))
         tensor([ 1.,  1., -1.], dtype=torch.float64)
     '''
-    out = torch.sign(input, out=None)
-    out[out==0] = 1
-    return out
+    return torch.sign(torch.sign(input) * 2 + 1)
 
 
 def cumops_(input, dim, ops):

pypose/lietensor/operation.py

@@ -211,9 +211,9 @@ def SE3_Adj(X):
 
 
 def SE3_Matrix(X):
-    T = torch.eye(4, device=X.device, dtype=X.dtype, requires_grad=False).repeat(X.shape[:-1]+(1, 1))
-    T[..., :3, :3] = SO3_Matrix(X[..., 3:])
-    T[..., :3, 3] = X[..., :3]
+    T = torch.cat([SO3_Matrix(X[..., 3:]), X[..., :3, None]], dim=-1)
+    E = torch.tensor([0, 0, 0, 1], dtype=T.dtype, device=T.device)
+    T = torch.cat([T, E.repeat(X.shape[:-1]+(1, 1))], dim=-2)
     return T
 
 
@@ -470,7 +470,7 @@ def backward(ctx, grad_output):
 
 
 class SO3_Act(torch.autograd.Function):
-
+    generate_vmap_rule = True
     @staticmethod
     def forward(X, p):
         Xv, Xw = X[..., :3], X[..., 3:]
@@ -498,7 +498,7 @@ def backward(ctx, grad_output):
 
 
 class SE3_Act(torch.autograd.Function):
-
+    generate_vmap_rule = True
     @staticmethod
     def forward(X, p):
         out = X[..., :3] + SO3_Act.apply(X[..., 3:], p)