Add helper function and implement AdvancedIncSubtensor

Author: HarshvirSandhu

pytensor/link/pytorch/dispatch/subtensor.py

@@ -1,5 +1,6 @@
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
 from pytensor.tensor.subtensor import (
+    AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
     AdvancedSubtensor,
     AdvancedSubtensor1,
@@ -10,18 +11,26 @@
 from pytensor.tensor.type_other import MakeSlice
 
 
+def check_negative_steps(index):
+    for i in index:
+        if isinstance(i, slice):
+            if i.step and isinstance(i.step, int) and i.step < 0:
+                raise NotImplementedError(
+                    "Negative step sizes are not supported in Pytorch"
+                )
+            elif i.step and not isinstance(i.step, int):
+                raise NotImplementedError(
+                    "Negative step sizes are not supported in Pytorch"
+                )
+
+
 @pytorch_funcify.register(Subtensor)
 def pytorch_funcify_Subtensor(op, node, **kwargs):
     idx_list = getattr(op, "idx_list", None)
 
     def subtensor(x, *ilists):
         indices = indices_from_subtensor(ilists, idx_list)
-        for i in indices:
-            if isinstance(i, slice):
-                if i.step and i.step < 0:
-                    raise NotImplementedError(
-                        "Negative step sizes are not supported in Pytorch"
-                    )
+        check_negative_steps(indices)
 
         return x[indices]
 
@@ -32,12 +41,7 @@ def subtensor(x, *ilists):
 @pytorch_funcify.register(AdvancedSubtensor)
 def pytorch_funcify_AdvSubtensor(op, node, **kwargs):
     def advsubtensor(x, *indices):
-        for i in indices:
-            if isinstance(i, slice):
-                if i.step and i.step < 0:
-                    raise NotImplementedError(
-                        "Negative step sizes are not supported in Pytorch"
-                    )
+        check_negative_steps(indices)
         return x[indices]
 
     return advsubtensor
@@ -52,21 +56,23 @@ def makeslice(*x):
 
 
 @pytorch_funcify.register(IncSubtensor)
-@pytorch_funcify.register(AdvancedIncSubtensor1)
 def pytorch_funcify_IncSubtensor(op, node, **kwargs):
     idx_list = getattr(op, "idx_list", None)
 
     if getattr(op, "set_instead_of_inc", False):
 
         def torch_fn(x, indices, y):
+            check_negative_steps(indices)
             x[indices] = y
             return x
 
     else:
 
         def torch_fn(x, indices, y):
-            x[indices] += y
-            return x
+            check_negative_steps(indices)
+            x1 = x.clone()
+            x1[indices] += y
+            return x1
 
     def incsubtensor(x, y, *ilist, torch_fn=torch_fn, idx_list=idx_list):
         indices = indices_from_subtensor(ilist, idx_list)
@@ -76,3 +82,30 @@ def incsubtensor(x, y, *ilist, torch_fn=torch_fn, idx_list=idx_list):
         return torch_fn(x, indices, y)
 
     return incsubtensor
+
+
+@pytorch_funcify.register(AdvancedIncSubtensor)
+@pytorch_funcify.register(AdvancedIncSubtensor1)
+def pytorch_funcify_AdvancedIncSubtensor(op, node, **kwargs):
+    if getattr(op, "set_instead_of_inc", False):
+
+        def torch_fn(x, indices, y):
+            check_negative_steps(indices)
+            x[indices] = y
+            return x
+
+    else:
+
+        def torch_fn(x, indices, y):
+            check_negative_steps(indices)
+            x1 = x.clone()
+            x1[indices] += y
+            return x1
+
+    def incsubtensor(x, y, *indices, torch_fn=torch_fn):
+        if len(indices) == 1:
+            indices = indices[0]
+
+        return torch_fn(x, indices, y)
+
+    return incsubtensor