[sparse] bcoo_extract incorrect for zero sparse dimensions
Found in the course of #13624 

Minimal repro:
```python
import jax.numpy as jnp
from jax.experimental import sparse

mat = jnp.zeros(3, dtype='float32')
indices = jnp.zeros((3, 2, 0), dtype='int32')

jax.jit(sparse.bcoo_extract)(indices, mat)  # Error
```
diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
index 247434523..81310b282 100644
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -400,7 +400,7 @@ def bcoo_extract(indices: Array, mat: Array) -> Array:
 @bcoo_extract_p.def_impl
 def _bcoo_extract_impl(indices, mat):
   mat = jnp.asarray(mat)
-  n_batch, n_sparse, _, _ = _validate_bcoo_indices(indices, mat.shape)
+  n_batch, n_sparse, _, nse = _validate_bcoo_indices(indices, mat.shape)
 
   ind_slices = tuple(np.zeros(s, int) if i_s == 1 else np.arange(s)
                      for s, i_s in zip(mat.shape[:n_batch], indices.shape[:n_batch]))
@@ -412,8 +412,13 @@ def _bcoo_extract_impl(indices, mat):
   batch_ind = tuple(grid)[:-1]
 
   if not sparse_ind + batch_ind:
-    return mat[None]
-  return mat.at[batch_ind + sparse_ind].get(mode='fill', fill_value=0)
+    result = mat[None]
+  else:
+    result = mat.at[batch_ind + sparse_ind].get(mode='fill', fill_value=0)
+  if n_sparse == 0 and nse != 1:
+    result = lax.broadcast_in_dim(
+      result, _tuple_replace(result.shape, n_batch, nse), range(result.ndim))
+  return result
 
 @bcoo_extract_p.def_abstract_eval
 def _bcoo_extract_abstract_eval(indices, mat):
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
index 7a9131082..27b029d0a 100644
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -904,6 +904,17 @@ class BCOOTest(sptu.SparseTestCase):
     self.assertEqual(j1.shape, data.shape + M.shape)
     self.assertEqual(hess.shape, data.shape + 2 * M.shape)
 
+  def test_bcoo_extract_zero_nse(self):
+    # Regression test for https://github.com/google/jax/issues/13653
+
+    # (n_batch, n_sparse, n_dense) = (1, 0, 0), nse = 2
+    args_maker = lambda: (jnp.zeros((3, 2, 0), dtype='int32'), jnp.arange(3))
+    self._CompileAndCheck(sparse.bcoo_extract, args_maker)
+
+    # (n_batch, n_sparse, n_dense) = (0, 0, 1), nse = 2
+    args_maker = lambda: (jnp.zeros((2, 0), dtype='int32'), jnp.arange(3))
+    self._CompileAndCheck(sparse.bcoo_extract, args_maker)
+
   @jtu.sample_product(
     [dict(shape=shape, n_batch=n_batch, n_dense=n_dense)
       for shape in [(), (5,), (5, 8), (3, 4, 5), (3, 4, 3, 2)]
['tests/sparse_test.py', 'jax/experimental/sparse/bcoo.py']