Repository: google/jax

Issue:
sparse.BCOO.fromdense cannot be called inside a function which is jittable
My use case is the following.  Inside a function `f`, after some computation, I get a matrix for which I am sure that it is a sparse matrix. I wish to convert this matrix to a sparse format using BCOO so that I can make the following computations efficient.  Ideally, I should be able to jit `f`. It doesn't look possible right now. 

Here is a simple example:

```python
from jax.experimental import sparse
jsum = sparse.sparsify(jnp.sum)

def test_bcoo(A):
    B  = sparse.BCOO.fromdense(A)
    return jsum(B)

print(test_bcoo(jnp.eye(4)))

test_bcoo_jit = jax.jit(test_bcoo)
print(test_bcoo_jit(jnp.eye(4)))
```

`test_bcoo` runs perfectly but `test_bcoo_jit` fails with following error: 

```
---------------------------------------------------------------------------
ConcretizationTypeError                   Traceback (most recent call last)
<ipython-input-91-4e719b763a0d> in <module>
      9 
     10 test_bcoo_jit = jax.jit(test_bcoo)
---> 11 print(test_bcoo_jit(jnp.eye(4)))

    [... skipping hidden 13 frame]

<ipython-input-91-4e719b763a0d> in test_bcoo(A)
      3 
      4 def test_bcoo(A):
----> 5     B  = sparse.BCOO.fromdense(A)
      6     return jsum(B)
      7 

/opt/anaconda3/envs/jax/lib/python3.9/site-packages/jax/experimental/sparse/bcoo.py in fromdense(cls, mat, nse, index_dtype, n_dense, n_batch)
    958   def fromdense(cls, mat, *, nse=None, index_dtype=np.int32, n_dense=0, n_batch=0):
    959     """Create a BCOO array from a (dense) :class:`DeviceArray`."""
--> 960     return cls(bcoo_fromdense(mat, nse=nse, index_dtype=index_dtype, n_dense=n_dense, n_batch=n_batch), shape=mat.shape)
    961 
    962   @classmethod

/opt/anaconda3/envs/jax/lib/python3.9/site-packages/jax/experimental/sparse/bcoo.py in bcoo_fromdense(mat, nse, n_batch, n_dense, index_dtype)
    216   if nse is None:
    217     nse = _bcoo_nse(mat, n_batch, n_dense)
--> 218   nse = core.concrete_or_error(operator.index, nse, "nse argument of bcoo_fromdense")
    219   return bcoo_fromdense_p.bind(mat, nse=nse, n_batch=n_batch, n_dense=n_dense,
    220                                index_dtype=index_dtype)

/opt/anaconda3/envs/jax/lib/python3.9/site-packages/jax/core.py in concrete_or_error(force, val, context)
   1009       return force(val.aval.val)
   1010     else:
-> 1011       raise ConcretizationTypeError(val, context)
   1012   else:
   1013     return force(val)

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=0/1)>
nse argument of bcoo_fromdense
While tracing the function test_bcoo at <ipython-input-91-4e719b763a0d>:4 for jit, this concrete value was not available in Python because it depends on the value of the argument 'A'.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError
``` 




Discussion:
Thanks for the report – we should improve this error, but it is expected that `sparse.BCOO.fromdense` will fail within JIT unless you pass a static value to `nse` (number of specified elements):
```python
def test_bcoo(A):
    B  = sparse.BCOO.fromdense(A, nse=4)
    return jsum(B)
print(test_bcoo(jnp.eye(4)))
# 4.0
print(jax.jit(test_bcoo)(jnp.eye(4)))
# 4.0
```
This is for the same reasons that simpler functions like `jnp.where`, or `jnp.nonzero`, or `jnp.unique`, will fail if used within JIT: JAX array shapes must be statically known, and not depend on the contents of traced arrays. It is a fundamental limitation of JAX's compilation model.

In general, I'd suggest treating sparse matrix creation as something that happens as a setup step before you call your workhorse function, so e.g. something like this:
```python
B  = sparse.BCOO.fromdense(A)

def test_bcoo(B):
    return jsum(B)
```
Then the number of specified elements can be determined automatically while the input matrix is still concrete.

Does that make sense?
#8305 improves this error message. It now says:
```
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>

The error arose for the nse argument of bcoo_fromdense. In order for BCOO.fromdense()
to be used in traced or compiled code, you must pass a concrete value to the nse
(number of specified elements) argument
```

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
@@ -197,6 +197,12 @@ def _bcoo_todense_batching_rule(batched_args, batch_dims, *, shape):
 bcoo_fromdense_p = core.Primitive('bcoo_fromdense')
 bcoo_fromdense_p.multiple_results = True
 
+_TRACED_NSE_ERROR = """
+The error arose for the nse argument of bcoo_fromdense. In order for BCOO.fromdense()
+to be used in traced/compiled code, you must pass a concrete value to the nse
+(number of specified elements) argument.
+"""
+
 def bcoo_fromdense(mat, *, nse=None, n_batch=0, n_dense=0, index_dtype=jnp.int32):
   """Create COO-format sparse matrix from a dense matrix.
 
@@ -215,7 +221,7 @@ def bcoo_fromdense(mat, *, nse=None, n_batch=0, n_dense=0, index_dtype=jnp.int32
   mat = jnp.asarray(mat)
   if nse is None:
     nse = _bcoo_nse(mat, n_batch, n_dense)
-  nse = core.concrete_or_error(operator.index, nse, "nse argument of bcoo_fromdense")
+  nse = core.concrete_or_error(operator.index, nse, _TRACED_NSE_ERROR)
   return bcoo_fromdense_p.bind(mat, nse=nse, n_batch=n_batch, n_dense=n_dense,
                                index_dtype=index_dtype)
 


Test patch: