Implemented Eye Op in PyTorch

- Added support for diagonal offset (param `k`)

Author: twaclaw

pytensor/link/pytorch/dispatch/basic.py

@@ -6,7 +6,7 @@
 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
+from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, Eye
 
 
 @singledispatch
@@ -89,3 +89,19 @@ def arange(start, stop, step):
         return torch.arange(start, stop, step, dtype=dtype)
 
     return arange
+
+
+@pytorch_funcify.register(Eye)
+def pytorch_funcify_eye(op, **kwargs):
+    dtype = getattr(torch, op.dtype)
+
+    def eye(N, M, k):
+        mjr, mnr = (M, N) if k > 0 else (N, M)
+        k_abs = abs(k)
+        zeros = torch.zeros(N, M, dtype=dtype)
+        if k_abs < mjr:
+            l_ones = min(mjr - k_abs, mnr)
+            return zeros.diagonal_scatter(torch.ones(l_ones, dtype=dtype), k)
+        return zeros
+
+    return eye

tests/link/pytorch/test_basic.py

@@ -12,7 +12,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
 from pytensor.raise_op import CheckAndRaise
-from pytensor.tensor import alloc, arange, as_tensor, empty
+from pytensor.tensor import alloc, arange, as_tensor, empty, eye
 from pytensor.tensor.type import scalar, vector
 
 
@@ -235,3 +235,20 @@ def test_arange():
         FunctionGraph([start, stop, step], [out]),
         [np.array(1), np.array(10), np.array(2)],
     )
+
+
+def test_eye():
+    N = scalar("N", dtype="int64")
+    M = scalar("M", dtype="int64")
+    k = scalar("k", dtype="int64")
+
+    out = eye(N, M, k, dtype="int16")
+
+    trange = range(1, 6)
+    for _N in trange:
+        for _M in trange:
+            for _k in list(range(_M + 2)) + [-x for x in range(1, _N + 2)]:
+                compare_pytorch_and_py(
+                    FunctionGraph([N, M, k], [out]),
+                    [np.array(_N + 1), np.array(_M + 1), np.array(_k)],
+                )