beignet.func.space

Genentech · Apr 22, 2024 · ba76031 · ba76031
1 parent 0c5b18d
commit ba76031
Showing 1 changed file with 48 additions and 4 deletions.
diff --git a/src/beignet/operators/_transform.py b/src/beignet/operators/_transform.py
@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import torch
 from torch import Tensor
 from torch.autograd import Function
 
@@ -47,7 +48,27 @@ def forward(transformation: Tensor, position: Tensor) -> Tensor:
             Tensor
                 Affine transformed position of shape `(..., dimension)`.
             """
-            return _transform(transformation, position)
+            indexes = "".join(
+                [chr(ord("a") + index) for index in range(position.ndim - 1)]
+            )
+
+            match transformation.ndim:
+                case 0:
+                    return position * transformation
+                case 1:
+                    return torch.einsum(
+                        f"i,{indexes}i->{indexes}i",
+                        transformation,
+                        position,
+                    )
+                case 2:
+                    return torch.einsum(
+                        f"ij,{indexes}j->{indexes}i",
+                        transformation,
+                        position,
+                    )
+                case _:
+                    raise ValueError
 
         @staticmethod
         def setup_context(ctx, inputs, output):
@@ -65,11 +86,34 @@ def jvp(
 
             output = _transform(transformation, position)
 
-            grad_output = grad_position + _transform(
-                grad_transformation,
-                position,
+            indexes = "".join(
+                [chr(ord("a") + index) for index in range(position.ndim - 1)]
             )
 
+            match grad_transformation.ndim:
+                case 0:
+                    transformed_grad = position * grad_transformation
+
+                    grad_output = grad_position + transformed_grad
+                case 1:
+                    transformed_grad = torch.einsum(
+                        f"i,{indexes}i->{indexes}i",
+                        grad_transformation,
+                        position,
+                    )
+
+                    grad_output = grad_position + transformed_grad
+                case 2:
+                    transformed_grad = torch.einsum(
+                        f"ij,{indexes}j->{indexes}i",
+                        grad_transformation,
+                        position,
+                    )
+
+                    grad_output = grad_position + transformed_grad
+                case _:
+                    raise ValueError
+
             return output, grad_output
 
         @staticmethod