Refactor tests for Eye Op in PyTorch

Author: twaclaw

tests/link/pytorch/test_basic.py

@@ -282,13 +282,13 @@ def test_eye():
     M = scalar("M", dtype="int64")
     k = scalar("k", dtype="int64")
 
-    out = eye(N, M, k, dtype="int16")
+    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 _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), np.array(_M), np.array(_k)],
-                )
+                np.testing.assert_array_equal(fn(_N, _M, _k), np.eye(_N, _M, _k))