Fix sizing of ones padded to position for homogenous coordinates

Author: LemonPi

src/pytorch_kinematics/transforms/transform3d.py

@@ -185,13 +185,13 @@ def __init__(
             self._matrix = matrix.view(-1, 4, 4)
 
         if pos is not None:
-            ones = torch.ones(1, dtype=dtype, device=device)
+            ones = torch.ones([1], dtype=dtype, device=device)
             if not torch.is_tensor(pos):
                 pos = torch.tensor(pos, dtype=dtype, device=device)
-            if pos.ndim in (2, 3) and pos.shape[0] > 1 and self._matrix.shape[0] == 1:
-                self._matrix = self._matrix.repeat(pos.shape[0], 1, 1)
-                ones = ones.repeat(pos.shape[0], 1)
-            # self._matrix[:, :3, 3] = pos
+            if pos.ndim in (2, 3):
+                ones = ones.repeat(*pos.shape[:-1], 1)
+                if pos.ndim in (2, 3) and pos.shape[0] > 1 and self._matrix.shape[0] == 1:
+                    self._matrix = self._matrix.repeat(pos.shape[0], 1, 1)
             pos_h = torch.cat((pos, ones), dim=-1).reshape(-1, 4, 1)
             self._matrix = torch.cat((self._matrix[:, :, :3], pos_h), dim=-1)
 

tests/test_jacobian.py

@@ -176,8 +176,8 @@ def get_pt(th):
     pk_end = timer()
     # we can only compare the positional parts
     assert torch.allclose(j1_, j2[:, :3], atol=1e-6)
-    print(f"for N={N} on {d} autograd:{(pk_start - autograd_start) * 1000}ms "
-          f"pytorch-kinematics:{(pk_end - pk_start) * 1000}ms")
+    print(f"for N={N} on {d} autograd:{(pk_start - autograd_start) * 1000}ms")
+    print(f"for N={N} on {d} pytorch-kinematics:{(pk_end - pk_start) * 1000}ms")
     # if we have functools (for pytorch>=1.13.0 it comes with installing pytorch)
     try:
         import functorch