Fix Subtensor and AdvSubtensor Ops

Author: HarshvirSandhu

pytensor/compile/mode.py

@@ -471,6 +471,7 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
             "BlasOpt",
             "fusion",
             "inplace",
+            "local_uint_constant_indices",
         ],
     ),
 )

pytensor/link/pytorch/dispatch/basic.py

@@ -1,29 +1,40 @@
 from functools import singledispatch
 from types import NoneType
 
+import numpy as np
 import torch
 
 from pytensor.compile.ops import DeepCopyOp
 from pytensor.graph.fg import FunctionGraph
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import CheckAndRaise
-from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, Eye, Join, MakeVector
+from pytensor.tensor.basic import (
+    Alloc,
+    AllocEmpty,
+    ARange,
+    Eye,
+    Join,
+    MakeVector,
+    TensorFromScalar,
+)
 
 
 @singledispatch
-def pytorch_typify(data, dtype=None, **kwargs):
-    r"""Convert instances of PyTensor `Type`\s to PyTorch types."""
-    if isinstance(data, NoneType):
-        raise NotImplementedError(f"pytorch_typify is not implemented for {type(data)}")
-    if isinstance(data, slice):
-        raise NotImplementedError(f"pytorch_typify is not implemented for {type(data)}")
-    else:
-        return torch.as_tensor(data, dtype=dtype)
+def pytorch_typify(data, **kwargs):
+    raise NotImplementedError(f"pytorch_typify is not implemented for {type(data)}")
 
 
+@pytorch_typify.register(np.ndarray)
+@pytorch_typify.register(torch.Tensor)
+def pytorch_typify_tensor(data, dtype=None, **kwargs):
+    return torch.as_tensor(data, dtype=dtype)
+
+
+@pytorch_typify.register(slice)
 @pytorch_typify.register(NoneType)
-def pytorch_typify_None(data, **kwargs):
-    return None
+@pytorch_typify.register(np.number)
+def pytorch_typify_scalar(data, **kwargs):
+    return data
 
 
 @singledispatch
@@ -137,3 +148,11 @@ def makevector(*x):
         return torch.tensor(x, dtype=torch_dtype)
 
     return makevector
+
+
+@pytorch_funcify.register(TensorFromScalar)
+def pytorch_funcify_TensorFromScalar(op, **kwargs):
+    def tensorfromscalar(x):
+        return x
+
+    return tensorfromscalar

pytensor/link/pytorch/dispatch/subtensor.py

@@ -7,34 +7,48 @@
     Subtensor,
     indices_from_subtensor,
 )
+from pytensor.tensor.type_other import MakeSlice
 
 
 @pytorch_funcify.register(Subtensor)
-@pytorch_funcify.register(AdvancedSubtensor)
-@pytorch_funcify.register(AdvancedSubtensor1)
 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)
-        new_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"
+                    )
+
+        return x[indices]
 
+    return subtensor
+
+
+@pytorch_funcify.register(AdvancedSubtensor1)
+@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"
                     )
-                new_indices.append(i)
-            else:
-                new_indices.append(i.tolist())
+        return x[indices]
 
-        if len(indices) == 1:
-            indices = indices[0]
+    return advsubtensor
 
-        return x[tuple(new_indices)]
 
-    return subtensor
+@pytorch_funcify.register(MakeSlice)
+def pytorch_funcify_makeslice(op, **kwargs):
+    def makeslice(*x):
+        return slice(x)
+
+    return makeslice
 
 
 @pytorch_funcify.register(IncSubtensor)
@@ -45,12 +59,14 @@ def pytorch_funcify_IncSubtensor(op, node, **kwargs):
     if getattr(op, "set_instead_of_inc", False):
 
         def torch_fn(x, indices, y):
-            return x.at[indices].set(y)
+            x[indices] = y
+            return x
 
     else:
 
         def torch_fn(x, indices, y):
-            return x.at[indices].add(y)
+            x[indices] += y
+            return x
 
     def incsubtensor(x, y, *ilist, torch_fn=torch_fn, idx_list=idx_list):
         indices = indices_from_subtensor(ilist, idx_list)

tests/link/pytorch/test_subtensor.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+import pytensor.scalar as ps
 import pytensor.tensor as pt
 from pytensor.configdefaults import config
 from pytensor.graph.fg import FunctionGraph
@@ -15,6 +16,7 @@ def test_pytorch_Subtensor():
 
     # Basic indices
     out_pt = x_pt[1, 2, 0]
+
     assert isinstance(out_pt.owner.op, pt_subtensor.Subtensor)
     out_fg = FunctionGraph([x_pt], [out_pt])
     compare_pytorch_and_py(out_fg, [x_np])
@@ -34,6 +36,27 @@ def test_pytorch_Subtensor():
     out_fg = FunctionGraph([x_pt], [out_pt])
     compare_pytorch_and_py(out_fg, [x_np])
 
+    a_pt = ps.int64("a")
+    a_np = 1
+    out_pt = x_pt[a_pt, 2, a_pt:2]
+
+    assert isinstance(out_pt.owner.op, pt_subtensor.Subtensor)
+    out_fg = FunctionGraph([x_pt, a_pt], [out_pt])
+    compare_pytorch_and_py(out_fg, [x_np, a_np])
+
+    # Negative step
+    with pytest.raises(
+        NotImplementedError, match="Negative step sizes are not supported in Pytorch"
+    ):
+        out_pt = x_pt[::-1]
+        out_fg = FunctionGraph([x_pt], [out_pt])
+        compare_pytorch_and_py(out_fg, [x_np])
+
+
+def test_pytorch_AdvSubtensor():
+    shape = (3, 4, 5)
+    x_pt = pt.tensor("x", shape=shape, dtype="int")
+    x_np = np.arange(np.prod(shape)).reshape(shape)
     # Advanced indexing
     out_pt = pt_subtensor.advanced_subtensor1(x_pt, [1, 2])
     assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor1)
@@ -51,24 +74,32 @@ def test_pytorch_Subtensor():
     out_fg = FunctionGraph([x_pt], [out_pt])
     compare_pytorch_and_py(out_fg, [x_np])
 
-    with pytest.raises(NotImplementedError):
-        out_pt = x_pt[[1, 2], :, [3, 4]]
-        assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor)
-        out_fg = FunctionGraph([x_pt], [out_pt])
-        compare_pytorch_and_py(out_fg, [x_np])
+    out_pt = x_pt[[1, 2], :, [3, 4]]
+    assert isinstance(out_pt.owner.op, pt_subtensor.AdvancedSubtensor)
+    out_fg = FunctionGraph([x_pt], [out_pt])
+    compare_pytorch_and_py(out_fg, [x_np])
 
-    # Flipping
-    with pytest.raises(
-        NotImplementedError, match="Negative step sizes are not supported in Pytorch"
-    ):
-        out_pt = x_pt[::-1]
-        out_fg = FunctionGraph([x_pt], [out_pt])
-        compare_pytorch_and_py(out_fg, [x_np])
+    out_pt = x_pt[[1, 2], None]
+    out_fg = FunctionGraph([x_pt], [out_pt])
+    compare_pytorch_and_py(out_fg, [x_np])
+
+    a_pt = ps.int64("a")
+    a_np = 2
+
+    out_pt = x_pt[[1, a_pt], a_pt]
+    out_fg = FunctionGraph([x_pt, a_pt], [out_pt])
+    compare_pytorch_and_py(out_fg, [x_np, a_np])
 
     out_pt = x_pt[np.random.binomial(1, 0.5, size=(3, 4, 5)).astype(bool)]
     out_fg = FunctionGraph([x_pt], [out_pt])
     compare_pytorch_and_py(out_fg, [x_np])
 
+    a_pt = pt.tensor3("a", dtype="bool")
+    a_np = np.random.binomial(1, 0.5, size=(3, 4, 5)).astype(bool)
+    out_pt = x_pt[a_pt]
+    out_fg = FunctionGraph([x_pt, a_pt], [out_pt])
+    compare_pytorch_and_py(out_fg, [x_np, a_np])
+
 
 def test_pytorch_IncSubtensor():
     rng = np.random.default_rng(42)