Improve repr for empty jax.Array

jax-ml · Feb 5, 2024 · d9cbd7b · d9cbd7b
1 parent e224c3d
commit d9cbd7b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/jax/_src/array.py b/jax/_src/array.py
@@ -376,8 +376,11 @@ def __repr__(self):
 
  if self.is_fully_addressable or self.is_fully_replicated:
  line_width = np.get_printoptions()["linewidth"]
- s = np.array2string(self._value, prefix=prefix, suffix=',',
- separator=', ', max_line_width=line_width)
+ if self.size == 0:
+ s = f"[], shape={self.shape}"
+ else:
+ s = np.array2string(self._value, prefix=prefix, suffix=',',
+ separator=', ', max_line_width=line_width)
  last_line_len = len(s) - s.rfind('\n') + 1
  sep = ' '
  if last_line_len + len(dtype_str) + 1 > line_width:

diff --git a/tests/array_test.py b/tests/array_test.py
@@ -231,6 +231,12 @@ def test_repr(self):
  input_shape, jax.sharding.NamedSharding(global_mesh, P('x', 'y')))
  self.assertStartsWith(repr(arr), "Array(")
 
+ def test_empty_repr(self):
+ shape = (0, 5)
+ dtype = 'float32'
+ x = jnp.empty(shape, dtype)
+ self.assertEqual(repr(x), f"Array([], shape={shape}, dtype={dtype})")
+
  def test_jnp_array(self):
  arr = jnp.array([1, 2, 3])
  self.assertIsInstance(arr, array.ArrayImpl)