Add torch support to @crazy decorator

src/pystencils_reco/_crazy_decorator.py

@@ -10,6 +10,7 @@
 
 import inspect
 
+from pystencils.autodiff.backends._pytorch import torch_dtype_to_numpy
 from pystencils.field import Field
 
 try:
@@ -23,13 +24,28 @@ except ImportError:
     tf = None
 
 
+def _create_field_from_array_like(field_name, maybe_array):
+
+    if torch:
+        # Torch tensors don't have t.strides but t.stride(dim). Let's fix that!
+        if isinstance(maybe_array, torch.Tensor):
+            fake_array = object()
+            fake_array.strides = [maybe_array.stride(i) for i in range(len(maybe_array.shape))]
+            fake_array.shape = maybe_array.shape
+            fake_array.dtype = torch_dtype_to_numpy(maybe_array.dtype)
+            field = Field.create_from_numpy_array(field_name, fake_array)
+            return field
+
+    return Field.create_from_numpy_array(field_name, maybe_array)
+
+
 def crazy(function):
 
     def wrapper(*args, **kwargs):
         arg_names = inspect.getfullargspec(function).args
-        compile_args = [Field.create_from_numpy_array(
+        compile_args = [_create_field_from_array_like(
             arg_names[i], a) if hasattr(a, '__array__') else a for i, a in enumerate(args)]
-        compile_kwargs = {k: Field.create_from_numpy_array(str(k), a) if hasattr(
+        compile_kwargs = {k: _create_field_from_array_like(str(k), a) if hasattr(
             a, '__array__') else a for (k, a) in kwargs.items()}
 
         assignments = function(*compile_args, **compile_kwargs)

tests/test_pytorch.py

@@ -7,6 +7,8 @@
 """
 
 """
+import torch
+
 import pystencils
 from pystencils.autodiff import torch_tensor_from_field
 from pystencils_reco.filters import mean_filter
@@ -43,7 +45,22 @@ def test_pytorch_gpu():
     print(torch_op)
 
 
+def test_pytorch_from_tensors():
+    block_stencil = BallStencil(1, ndim=2)
+
+    x, y = pystencils.fields('x,y: float32[100,100]')
+    x_tensor = torch_tensor_from_field(x, requires_grad=True, cuda=True)
+    y_tensor = torch_tensor_from_field(y, cuda=True)
+    filter = mean_filter(x_tensor, y_tensor, block_stencil)
+    print(filter)
+    print(filter.backward())
+
+    torch_op = filter.create_pytorch_op(x=x_tensor+1, y=y_tensor)
+    print(torch_op)
+
+
 def main():
+    # test_pytorch()
     test_pytorch()