Repository: google/jax

Issue:
Bug in np.eye.
As suggested, it seems like np.eye broke at some point. To repro try, `np.allclose(onp.eye(N), np.eye(N))`. 

Upon closer inspection, it seems like this only kicks in at large N. Try e.g. 5000.


Discussion:
Interestingly, it seems that `onp.allclose` works and `np.allclose` doesn't...

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -3646,7 +3646,7 @@ class _IotaConstant(xla.DeviceConstant):
 
   def __init__(self, dtype, shape, axis):
     self.shape = shape
-    self.dtype = dtype
+    self.dtype = onp.dtype(dtype)
     self.ndim = len(shape)
     self.size = prod(shape)
     self._npy_value = None
@@ -3675,7 +3675,7 @@ class _EyeConstant(xla.DeviceConstant):
 
   def __init__(self, shape, axes, dtype):
     self.shape = shape
-    self.dtype = dtype
+    self.dtype = onp.dtype(dtype)
     self.ndim = len(shape)
     self.size = prod(shape)
     self._npy_value = None
@@ -3700,7 +3700,7 @@ def constant_handler(c, diag_const, canonicalize_types=True):
     else:
       etype = xla_bridge.dtype_to_etype_exact(diag_const.dtype)
     etype = xla_bridge.dtype_to_etype(diag_const.dtype)
-    iotas = [c.BroadcastedIota(onp.bool_, diag_const.shape, axis)
+    iotas = [c.BroadcastedIota(onp.uint32, diag_const.shape, axis)
              for axis in diag_const.axes]
     eyes = [c.Eq(i1, i2) for i1, i2 in zip(iotas[:-1], iotas[1:])]
     return c.ConvertElementType(_reduce(c.And, eyes), etype)


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1481,6 +1481,10 @@ def testArange(self):
     self.assertFalse(type(lnp.arange(77)) == type(onp.arange(77)))
     self.assertTrue(type(lnp.arange(77)) == type(lax.iota(onp.int32, 77)))
 
+  def testIssue728(self):
+    assert lnp.allclose(lnp.eye(5000), onp.eye(5000))
+    self.assertEqual(0, onp.sum(lnp.eye(1050) - onp.eye(1050)))
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1343,13 +1343,16 @@ def _CheckDeviceConstant(self, make_const, expected):
     self.assertAllClose(argument_result, expected, check_dtypes=True)
     self.assertAllClose(jit_result, expected, check_dtypes=True)
 
+    # ensure repr doesn't crash
+    repr(make_const())
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_fill={}".format(
           jtu.format_shape_dtype_string(shape, dtype) if dtype else shape,
           fill_value),
        "shape": shape, "dtype": dtype, "fill_value": fill_value}
       for dtype in itertools.chain(default_dtypes, [None])
-      for shape in [(), (3,), (2, 3), (2, 3, 4)]
+      for shape in [(), (3,), (2, 3), (2, 3, 4), (1001, 1001)]
       for fill_value in [0, 1, onp.pi]))
   def testFilledConstant(self, shape, fill_value, dtype):
     make_const = lambda: lax.full(shape, fill_value, dtype)
@@ -1361,7 +1364,7 @@ def testFilledConstant(self, shape, fill_value, dtype):
           jtu.format_shape_dtype_string(shape, dtype), dimension),
        "shape": shape, "dtype": dtype, "dimension": dimension}
       for dtype in default_dtypes
-      for shape in [(), (3,), (2, 3), (2, 3, 4)]
+      for shape in [(), (3,), (2, 3), (2, 3, 4), (1001, 1001), (101, 101, 101)]
       for dimension in range(len(shape))))
   def testIotaConstant(self, dtype, shape, dimension):
     make_const = lambda: lax.broadcasted_iota(dtype, shape, dimension)
@@ -1386,6 +1389,7 @@ def testIotaConstant(self, dtype, shape, dimension):
           [(2, 3, 4), (0, 1, 2)],
           [(2, 3, 4, 2), (0, 1, 2)],
           [(2, 3, 4, 2), (0, 2, 3)],
+          [(1001, 1001), (0, 1)],
       ]))
   def testEyeConstant(self, dtype, shape, axes):
     make_const = lambda: lax.broadcasted_eye(dtype, shape, axes)