Repository: ivy-llc/ivy

Issue:
grad_fn
Implement torch.Tensor.grad_fn


Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/torch/func_wrapper.py b/ivy/functional/frontends/torch/func_wrapper.py
--- a/ivy/functional/frontends/torch/func_wrapper.py
+++ b/ivy/functional/frontends/torch/func_wrapper.py
@@ -7,13 +7,90 @@
 import ivy.functional.frontends.torch as torch_frontend
 
 
-def _from_ivy_array_to_torch_frontend_tensor(x, nested=False, include_derived=None):
+class AccumulateGrad:
+    def __init__(self) -> None:
+        self.next_functions = ()
+        self.__name__ = "AccumulateGrad"
+
+    def __repr__(self):
+        return self.__name__
+
+    def __eq__(self, __value: object) -> bool:
+        return self.__name__ == __value
+
+    def __call__(self, grads):
+        for i in range(self.__self__.ndim):
+            grads = grads.sum(-1)
+        self.__self__._grads = grads
+        return None
+
+
+class GradFn:
+    def __init__(self, fn, inputs) -> None:
+        self._inputs = []
+        self._fns = []
+        self.next_functions = []
+        assert len(inputs) <= 2
+        if len(inputs) == 1 and isinstance(inputs[0], torch_frontend.Tensor):
+            self._inputs.append(inputs[0].detach())
+            d_fn = lambda x: fn(x)
+            self._fns.append(to_ivy_arrays_and_back(ivy.jac(d_fn)))
+            if inputs[0].grad_fn is not None:
+                self.next_functions.append(inputs[0].grad_fn)
+            elif inputs[0].requires_grad and inputs[0].is_leaf:
+                acc_grad = AccumulateGrad()
+                acc_grad.__self__ = inputs[0]
+                self.next_functions.append(acc_grad)
+        elif len(inputs) == 2:
+            if isinstance(inputs[0], torch_frontend.Tensor):
+                self._inputs.append(inputs[0].detach())
+                d_fn = lambda x: fn(x, inputs[1])
+                self._fns.append(to_ivy_arrays_and_back(ivy.jac(d_fn)))
+                if inputs[0].grad_fn is not None:
+                    self.next_functions.append(inputs[0].grad_fn)
+                elif inputs[0].requires_grad and inputs[0].is_leaf:
+                    acc_grad = AccumulateGrad()
+                    acc_grad.__self__ = inputs[0]
+                    self.next_functions.append(acc_grad)
+            if isinstance(inputs[1], torch_frontend.Tensor):
+                self._inputs.append(inputs[1].detach())
+                d_fn = lambda x: fn(inputs[0], x)
+                self._fns.append(to_ivy_arrays_and_back(ivy.jac(d_fn)))
+                if inputs[1].grad_fn is not None:
+                    self.next_functions.append(inputs[1].grad_fn)
+                elif inputs[1].requires_grad and inputs[1].is_leaf:
+                    acc_grad = AccumulateGrad()
+                    acc_grad.__self__ = inputs[1]
+                    self.next_functions.append(acc_grad)
+        self.__name__ = fn.__name__.capitalize() + "Backward"
+
+    def __call__(self, prev_grads):
+        return [
+            jac_fn(input_tensor) * prev_grads
+            for input_tensor, jac_fn in zip(self._inputs, self._fns)
+        ]
+
+    def __repr__(self):
+        return self.__name__
+
+    def __eq__(self, __value: object) -> bool:
+        return self.__name__ == __value
+
+
+def _from_ivy_array_to_torch_frontend_tensor(
+    x, nested=False, include_derived=None, requires_grad=False
+):
     if nested:
         return ivy.nested_map(
-            x, _from_ivy_array_to_torch_frontend_tensor, include_derived, shallow=False
+            x,
+            functools.partial(
+                _from_ivy_array_to_torch_frontend_tensor, requires_grad=requires_grad
+            ),
+            include_derived,
+            shallow=False,
         )
     elif isinstance(x, ivy.Array) or ivy.is_native_array(x):
-        a = torch_frontend.Tensor(x, _init_overload=True)
+        a = torch_frontend.Tensor(x, _init_overload=True, requires_grad=requires_grad)
         return a
     return x
 
@@ -91,7 +168,18 @@ def outputs_to_frontend_arrays_torch(*args, **kwargs):
                 ivy.unset_default_float_dtype()
         # convert all arrays in the return to `torch_frontend.Tensor` instances
         ret = _from_ivy_array_to_torch_frontend_tensor(
-            ret, nested=True, include_derived={tuple: True}
+            ret,
+            nested=True,
+            include_derived={tuple: True},
+            requires_grad=kwargs.get(
+                "requires_grad",
+                any(
+                    [
+                        isinstance(i, torch_frontend.Tensor) and i.requires_grad
+                        for i in args
+                    ]
+                ),
+            ),
         )
         array_fn = lambda x: ivy.is_array(x) or hasattr(x, "ivy_array")
         if "inplace" in kwargs and kwargs["inplace"]:
@@ -104,9 +192,32 @@ def outputs_to_frontend_arrays_torch(*args, **kwargs):
                 first_array._data = ret.ivy_array._data
             else:
                 first_array.ivy_array._data = ret.ivy_array._data
-            return first_array
-        else:
-            return ret
+            ret = first_array
+
+        # logic for setting is_leaf
+        if ret is not None and isinstance(ret, torch_frontend.Tensor):
+            if fn.__name__ in dir(torch_frontend.creation_ops):
+                ret.is_leaf = True
+            elif all(
+                [
+                    not isinstance(i, torch_frontend.Tensor)
+                    or (not i.requires_grad and not i.grad_fn)
+                    for i in args
+                ]
+            ):
+                ret.is_leaf = True
+            else:
+                ret.is_leaf = False
+        # set grad_fn
+        if any(
+            [isinstance(i, torch_frontend.Tensor) and i.requires_grad for i in args]
+        ):
+            # ToDo: Implement for unbind
+            grad_fn = GradFn(fn, args)
+            grad_fn.__self__ = ret
+            ret.grad_fn = grad_fn
+
+        return ret
 
     return outputs_to_frontend_arrays_torch
 
diff --git a/ivy/functional/frontends/torch/tensor.py b/ivy/functional/frontends/torch/tensor.py
--- a/ivy/functional/frontends/torch/tensor.py
+++ b/ivy/functional/frontends/torch/tensor.py
@@ -18,7 +18,7 @@
 
 
 class Tensor:
-    def __init__(self, array, device=None, _init_overload=False):
+    def __init__(self, array, device=None, _init_overload=False, requires_grad=False):
         if _init_overload:
             self._ivy_array = (
                 ivy.array(array) if not isinstance(array, ivy.Array) else array
@@ -29,6 +29,13 @@ def __init__(self, array, device=None, _init_overload=False):
                 array, dtype=torch_frontend.float32, device=device
             )
         self._grads = None
+        self._requires_grad = requires_grad
+        self.grad_fn = None
+        self._grads = None
+        if not _init_overload:
+            self._is_leaf = True
+        else:
+            self._is_leaf = False
 
     def __len__(self):
         return len(self._ivy_array)
@@ -85,6 +92,14 @@ def data(self):
     def grad(self):
         return self._grads
 
+    @property
+    def requires_grad(self):
+        return self._requires_grad
+
+    @property
+    def is_leaf(self):
+        return self._is_leaf
+
     # Setters #
     # --------#
 
@@ -94,6 +109,14 @@ def ivy_array(self, array):
             ivy.array(array) if not isinstance(array, ivy.Array) else array
         )
 
+    @requires_grad.setter
+    def requires_grad(self, requires_grad):
+        self._requires_grad = requires_grad
+
+    @is_leaf.setter
+    def is_leaf(self, is_leaf):
+        self._is_leaf = is_leaf
+
     # Instance Methods #
     # ---------------- #
     def reshape(self, *args, shape=None):


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_tensor.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_tensor.py
--- a/ivy_tests/test_ivy/test_frontends/test_torch/test_tensor.py
+++ b/ivy_tests/test_ivy/test_frontends/test_torch/test_tensor.py
@@ -190,6 +190,49 @@ def test_torch_tensor_property_grad():
     assert ivy.array_equal(x.grad, grads)
 
 
+@given(
+    dtype_x=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("valid", prune_function=False),
+    ),
+    requires_grad=st.booleans(),
+)
+def test_torch_tensor_property_requires_grad(
+    dtype_x,
+    requires_grad,
+):
+    _, data = dtype_x
+    x = Tensor(data[0], requires_grad=requires_grad)
+    ivy.utils.assertions.check_equal(x.requires_grad, requires_grad, as_array=False)
+    x.requires_grad = not requires_grad
+    ivy.utils.assertions.check_equal(x.requires_grad, not requires_grad, as_array=False)
+
+
+@given(
+    requires_grad=st.booleans(),
+)
+def test_torch_tensor_property_is_leaf(
+    requires_grad,
+):
+    x = Tensor(ivy.array([3.0]), requires_grad=requires_grad)
+    ivy.utils.assertions.check_equal(x.is_leaf, True, as_array=False)
+    y = x.pow(2)
+    ivy.utils.assertions.check_equal(y.is_leaf, not requires_grad, as_array=False)
+    z = y.detach()
+    ivy.utils.assertions.check_equal(z.is_leaf, True, as_array=False)
+
+
+def test_torch_tensor_property_grad_fn():
+    x = Tensor(ivy.array([3.0]), requires_grad=True)
+    ivy.utils.assertions.check_equal(x.grad_fn, None, as_array=False)
+    y = x.pow(2)
+    ivy.utils.assertions.check_equal(y.grad_fn, "PowBackward", as_array=False)
+    ivy.utils.assertions.check_equal(
+        y.grad_fn.next_functions[0], "AccumulateGrad", as_array=False
+    )
+    z = y.detach()
+    ivy.utils.assertions.check_equal(z.grad_fn, None, as_array=False)
+
+
 # chunk
 @pytest.mark.skip("Testing takes a lot of time")
 @handle_frontend_method(