Implemented Repeat and Unique Ops in PyTorch

Author: twaclaw

pytensor/link/pytorch/dispatch/extra_ops.py

@@ -1,7 +1,7 @@
 import torch
 
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
-from pytensor.tensor.extra_ops import CumOp
+from pytensor.tensor.extra_ops import CumOp, Repeat, Unique
 
 
 @pytorch_funcify.register(CumOp)
@@ -21,3 +21,44 @@ def cumop(x):
             return torch.cumprod(x, dim=dim)
 
     return cumop
+
+
+@pytorch_funcify.register(Repeat)
+def pytorch_funcify_Repeat(op, **kwargs):
+    axis = op.axis
+
+    def repeat(x, repeats, axis=axis):
+        return x.repeat_interleave(repeats, dim=axis)
+
+    return repeat
+
+
+@pytorch_funcify.register(Unique)
+def pytorch_funcify_Unique(op, **kwargs):
+    return_index = op.return_index
+
+    if return_index:
+        # TODO: evaluate whether is worth implementing this param
+        # (see https://github.com/pytorch/pytorch/issues/36748)
+        raise NotImplementedError("return_index is not implemented for pytorch")
+
+    axis = op.axis
+    return_inverse = op.return_inverse
+    return_counts = op.return_counts
+
+    def unique(
+        x,
+        return_index=return_index,
+        return_inverse=return_inverse,
+        return_counts=return_counts,
+        axis=axis,
+    ):
+        return torch.unique(
+            x,
+            sorted=True,
+            return_inverse=return_inverse,
+            return_counts=return_counts,
+            dim=axis,
+        )
+
+    return unique

tests/link/pytorch/test_extra_ops.py

@@ -41,3 +41,49 @@ def test_pytorch_CumOp(axis, dtype):
         out = pt.cumprod(a, axis=axis)
         fgraph = FunctionGraph([a], [out])
         compare_pytorch_and_py(fgraph, [test_value])
+
+
+def test_pytorch_Repeat():
+    a = pt.matrix("a", dtype="float64")
+
+    test_value = np.arange(6, dtype="float64").reshape((3, 2))
+
+    # Test along axis 0
+    out = pt.repeat(a, (1, 2, 3), axis=0)
+    fgraph = FunctionGraph([a], [out])
+    compare_pytorch_and_py(fgraph, [test_value])
+
+    # Test along axis 1
+    out = pt.repeat(a, (3, 3), axis=1)
+    fgraph = FunctionGraph([a], [out])
+    compare_pytorch_and_py(fgraph, [test_value])
+
+
+def test_pytorch_Unique():
+    a = pt.matrix("a", dtype="float64")
+
+    test_value = np.array(
+        [[1.0, 1.0, 2.0], [1.0, 1.0, 2.0], [3.0, 3.0, 0.0]], dtype="float64"
+    )
+
+    # Test along axis 0
+    out = pt.unique(a, axis=0)
+    fgraph = FunctionGraph([a], [out])
+    compare_pytorch_and_py(fgraph, [test_value])
+
+    # Test along axis 1
+    out = pt.unique(a, axis=1)
+    fgraph = FunctionGraph([a], [out])
+    compare_pytorch_and_py(fgraph, [test_value])
+
+    # Test with params
+    out = pt.unique(a, return_inverse=True, return_counts=True, axis=0)
+    fgraph = FunctionGraph([a], [out[0]])
+    compare_pytorch_and_py(fgraph, [test_value])
+
+    # Test with return_index=True
+    out = pt.unique(a, return_index=True, axis=0)
+    fgraph = FunctionGraph([a], [out[0]])
+
+    with pytest.raises(NotImplementedError):
+        compare_pytorch_and_py(fgraph, [test_value])