beignet.func.space

Genentech · Apr 22, 2024 · eaca190 · eaca190
1 parent f3f894b
commit eaca190
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/tests/beignet/operators/test__transform.py b/tests/beignet/operators/test__transform.py
@@ -1,6 +1,7 @@
 import beignet.operators
 import torch.func
 import torch.testing
+from torch import Tensor
 
 
 def test_transform():
@@ -20,11 +21,11 @@ def test_transform():
         ),
     )
 
-    def f(R):
-        return torch.sum(input**2)
+    def f(r: Tensor) -> Tensor:
+        return torch.sum(r**2)
 
-    def g(T, R):
-        return torch.sum(beignet.operators.transform(transformation, input) ** 2)
+    def g(t: Tensor, r: Tensor) -> Tensor:
+        return torch.sum(beignet.operators.transform(t, r) ** 2)
 
     torch.testing.assert_allclose(
         torch.func.grad(f)(beignet.operators.transform(transformation, input)),