Repository: google/jax

Issue:
`BCOO.fromdense` is not compatible with `jax.vmap`
### Description

Something unrelated that I bumped into whilst investigating #18244:
```python
import jax
import jax.numpy as jnp
from jax.experimental import sparse

matrix = jax.vmap(sparse.BCOO.fromdense)(jnp.arange(16.).reshape(1, 4, 4))
#   File ".../jax/experimental/sparse/util.py", line 106, in _count_stored_elements
#     return int(_count_stored_elements_per_batch(mat, n_batch, n_dense).max(initial=0))
#            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# jax.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with # shape int32[].
# The problem arose with the `int` function. If trying to convert the data type of a value, try using `x.astype(int)` or # `jnp.array(x, int)` instead.
```
I know about the `fromdense(..., n_batch=...)` argment, but I think it'd be reasonable for `fromdense` to occur within traced code.

### What jax/jaxlib version are you using?

JAX 0.4.19

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
This is working as intended. You need to pass a static value to `nse` in order to use `BCOO.fromdense` within `vmap` and other JAX transformations, because otherwise the output arrays have data-dependent size.
Ah, gotcha. In that case I think treat this as a report that the error message could be improved.
Thanks - I think the intent is that it would hit this line and raise a useful error, but it's clearly not doing that: https://github.com/google/jax/blob/20e583834ee8e8dc6e4d0c43d2eb86cfe9428f58/jax/experimental/sparse/bcoo.py#L289

This is the intended error message: https://github.com/google/jax/blob/20e583834ee8e8dc6e4d0c43d2eb86cfe9428f58/jax/experimental/sparse/bcoo.py#L244-L248

Resolving patch:
diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -20,7 +20,7 @@
 from functools import partial
 import math
 import operator
-from typing import Any, NamedTuple, Optional, Protocol, Union
+from typing import Any, NamedTuple, Protocol
 import warnings
 
 import numpy as np
@@ -262,9 +262,10 @@ def bcoo_fromdense(mat: Array, *, nse: int | None = None, n_batch: int = 0,
     mat_bcoo: BCOO representation of the matrix.
   """
   mat = jnp.asarray(mat)
-  if nse is None:
-    nse = _count_stored_elements(mat, n_batch, n_dense)
-  nse_int = core.concrete_or_error(operator.index, nse, _TRACED_NSE_ERROR)
+  nse_arr: int | Array | None = nse
+  if nse_arr is None:
+    nse_arr = _count_stored_elements(mat, n_batch, n_dense)
+  nse_int = core.concrete_or_error(operator.index, nse_arr, _TRACED_NSE_ERROR)
   return BCOO(_bcoo_fromdense(mat, nse=nse_int, n_batch=n_batch, n_dense=n_dense,
                               index_dtype=index_dtype),
               shape=mat.shape, indices_sorted=True, unique_indices=True)
diff --git a/jax/experimental/sparse/bcsr.py b/jax/experimental/sparse/bcsr.py
--- a/jax/experimental/sparse/bcsr.py
+++ b/jax/experimental/sparse/bcsr.py
@@ -194,9 +194,10 @@ def bcsr_fromdense(mat: ArrayLike, *, nse: int | None = None, n_batch: int = 0,
     mat_bcsr: BCSR representation of the matrix.
   """
   mat_array = jnp.asarray(mat)
-  if nse is None:
-    nse = _count_stored_elements(mat_array, n_batch, n_dense)
-  nse_int: int = core.concrete_or_error(operator.index, nse, _TRACED_NSE_ERROR)
+  nse_arr: int | Array | None = nse
+  if nse_arr is None:
+    nse_arr = _count_stored_elements(mat_array, n_batch, n_dense)
+  nse_int: int = core.concrete_or_error(operator.index, nse_arr, _TRACED_NSE_ERROR)
   return BCSR(_bcsr_fromdense(mat_array, nse=nse_int, n_batch=n_batch,
                               n_dense=n_dense, index_dtype=index_dtype),
               shape=mat_array.shape)
diff --git a/jax/experimental/sparse/util.py b/jax/experimental/sparse/util.py
--- a/jax/experimental/sparse/util.py
+++ b/jax/experimental/sparse/util.py
@@ -101,9 +101,9 @@ def _count_stored_elements_per_batch(mat: Array, n_batch: int = 0, n_dense: int
   mask = mask.sum(tuple(range(n_batch, mask.ndim)))
   return mask
 
-def _count_stored_elements(mat: Array, n_batch: int = 0, n_dense: int = 0) -> int:
+def _count_stored_elements(mat: Array, n_batch: int = 0, n_dense: int = 0) -> Array:
   """Return the number of stored elements (nse) of the given dense matrix."""
-  return int(_count_stored_elements_per_batch(mat, n_batch, n_dense).max(initial=0))
+  return _count_stored_elements_per_batch(mat, n_batch, n_dense).max(initial=0)
 
 def _dot_general_validated_shape(
     lhs_shape: tuple[int, ...], rhs_shape: tuple[int, ...],


Test patch: