Renamed vars in implemeantion of Eye Op in PyTorch

Author: twaclaw

pytensor/link/pytorch/dispatch/basic.py

@@ -104,15 +104,15 @@ def join(axis, *tensors):
 
 @pytorch_funcify.register(Eye)
 def pytorch_funcify_eye(op, **kwargs):
-    dtype = getattr(torch, op.dtype)
+    torch_dtype = getattr(torch, op.dtype)
 
     def eye(N, M, k):
-        mjr, mnr = (M, N) if k > 0 else (N, M)
+        major, minor = (M, N) if k > 0 else (N, M)
         k_abs = torch.abs(k)
-        zeros = torch.zeros(N, M, dtype=dtype)
-        if k_abs < mjr:
-            l_ones = torch.min(mjr - k_abs, mnr)
-            return zeros.diagonal_scatter(torch.ones(l_ones, dtype=dtype), k)
+        zeros = torch.zeros(N, M, dtype=torch_dtype)
+        if k_abs < major:
+            l_ones = torch.min(major - k_abs, minor)
+            return zeros.diagonal_scatter(torch.ones(l_ones, dtype=torch_dtype), k)
         return zeros
 
     return eye

tests/link/pytorch/test_basic.py

@@ -284,11 +284,9 @@ def test_eye():
 
     out = eye(N, M, k, dtype="float32")
 
-    trange = range(1, 6)
-
     fn = function([N, M, k], out, mode=pytorch_mode)
 
-    for _N in trange:
-        for _M in trange:
+    for _N in range(1, 6):
+        for _M in range(1, 6):
             for _k in list(range(_M + 2)) + [-x for x in range(1, _N + 2)]:
                 np.testing.assert_array_equal(fn(_N, _M, _k), np.eye(_N, _M, _k))