Repository: google/jax

Issue:
`repr` error on empty BCOO sparse matrix
I tried to initialize an empty BCOO matrix by:
```
from jax.experimental import sparse as jsparse
import jax.numpy as jnp

M = jsparse.BCOO((jnp.array([]), jnp.array([])), shape=(100, 100))
```

Just using the repr on `M` gives the following error: `<repr(<jax.experimental.sparse.ops.BCOO at 0x7f15fb4c9590>) failed: IndexError: tuple index out of range>`, which is unexpected.

On the other hand, initializing `M` with

```M = jsparse.BCOO((jnp.zeros(0, dtype='float32'), jnp.zeros((0, 2), dtype='int32')), shape=(100, 100)) ```

gives the correct repr output
```
M
BCOO(float32[100, 100], nse=0)
```




Discussion:
Thanks for the report!

I think we need some input validation on `__init__` – this is a bit difficult, though, because jax transforms sometimes instantiate pytrees with placeholder objects... so we'd have to somehow skip the validation in those cases.
I've also noted the following strange behavior:

```python
q = jsparse.BCOO.fromdense(jnp.zeros(100, dtype=jnp.float32))
q.data, q.indices
# (DeviceArray([], dtype=float32), DeviceArray([], dtype=int32))

jsparse.BCOO((jnp.array([], dtype=jnp.float32), jnp.array([], dtype=jnp.int32)), shape=(100,))
# <repr(<jax.experimental.sparse.ops.BCOO at 0x7f4337d49950>) failed: IndexError: tuple index out of range>
```

while 
```python
jsparse.BCOO((q.data, q.indices), shape=(100,))
# BCOO(float32[100], nse=0)
```

I think this is more of a bug on the `repr` side. Inspecting the shapes manually gives:

```python
q.data.shape, q.indices.shape
# (0,), (0,1)
```

Initializing the BCOO accordingly gives the expected output
```python
q = jsparse.BCOO((jnp.zeros(0), jnp.zeros((0,1))), shape=(100,))
q
# BCOO(float32[100], nse=0)
```
but this cannot be guessed by simply printing `q.data, q.indices`, as maybe it should?

Yeah, `indices` need to be of shape `[nse, rank]`. You can see this is the case, despite the empty array repr not being clear on this:
```python
import jax.experimental.sparse as jsparse
import jax.numpy as jnp

q = jsparse.BCOO.fromdense(jnp.zeros(100, dtype=jnp.float32))
q.indices.shape
# [0, 1]
```
As with the original error, if the indices are the wrong shape it will lead to issues.
And yes, I think that instantiating `BCOO` with an incorrectly-shaped index should result in an immediate error. But as mentioned above this is difficult because of how JAX transforms use placeholders within pytrees.

Resolving patch:
diff --git a/jax/experimental/sparse/ops.py b/jax/experimental/sparse/ops.py
--- a/jax/experimental/sparse/ops.py
+++ b/jax/experimental/sparse/ops.py
@@ -1394,7 +1394,15 @@ def __init__(self, args, *, shape):
     self.shape = shape
 
   def __repr__(self):
-    repr_ = f"{self.__class__.__name__}({self.dtype}{list(self.shape)}, nse={self.nse})"
+    name = self.__class__.__name__
+    try:
+      nse = self.nse
+      dtype = self.dtype
+      shape = list(self.shape)
+    except:
+      repr_ = f"{name}(<invalid>)"
+    else:
+      repr_ = f"{name}({dtype}{shape}, nse={nse})"
     if isinstance(self.data, core.Tracer):
       repr_ = f"{type(self.data).__name__}[{repr_}]"
     return repr_


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -1113,6 +1113,13 @@ def f(X, y):
 
 
 class SparseObjectTest(jtu.JaxTestCase):
+  def test_repr(self):
+    M = sparse.BCOO.fromdense(jnp.arange(5, dtype='float32'))
+    assert repr(M) == "BCOO(float32[5], nse=4)"
+
+    M_invalid = sparse.BCOO(([], []), shape=100)
+    assert repr(M_invalid) == "BCOO(<invalid>)"
+
   @parameterized.named_parameters(
     {"testcase_name": "_{}".format(Obj.__name__), "Obj": Obj}
     for Obj in [sparse.CSR, sparse.CSC, sparse.COO, sparse.BCOO])