Repository: google/jax

Issue:
sparse.eye calls Numpy preventing JIT
### Description

```
from jax import jit, numpy as jnp
from jax.experimental.sparse import eye


def f(A):
    n = A.shape[0]
    return eye(n)


A = jnp.ones([3, 3])
f(A)
# BCOO(float32[3, 3], nse=3)
jit(f)(A)
# TracerArrayConversionError ...
```

The problem function is:

https://github.com/google/jax/blob/42e367af9c58cb67496e235eddf325aa285600e4/jax/_src/lax/lax.py#L4708-L4713

I think the return could just be changed to `jnp.array`. 

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

v0.3.23

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

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report - looks like the implementation is a bit sloppy with respect to static and dynamic values. Who wrote this code anyway? [oh no...](https://github.com/google/jax/pull/10782) 😀 

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
@@ -2468,7 +2468,7 @@ def _eye(cls, N, M, k, *, dtype=None, index_dtype='int32', n_batch=0, n_dense=0)
       return cls.fromdense(jnp.eye(N, M, k, dtype=dtype),
                            n_batch=n_batch, n_dense=n_dense,
                            index_dtype=index_dtype)
-    k = jnp.asarray(k)
+
     if n_batch == 0:
       data = jnp.ones(diag_size, dtype=dtype)
       idx = jnp.arange(diag_size, dtype=index_dtype)
diff --git a/jax/experimental/sparse/coo.py b/jax/experimental/sparse/coo.py
--- a/jax/experimental/sparse/coo.py
+++ b/jax/experimental/sparse/coo.py
@@ -110,7 +110,6 @@ def _eye(cls, N, M, k, *, dtype=None, index_dtype='int32'):
       # if k is out of range, return an empty matrix.
       return cls._empty((N, M), dtype=dtype, index_dtype=index_dtype)
 
-    k = jnp.asarray(k)
     data = jnp.ones(diag_size, dtype=dtype)
     idx = jnp.arange(diag_size, dtype=index_dtype)
     zero = _const(idx, 0)
diff --git a/jax/experimental/sparse/csr.py b/jax/experimental/sparse/csr.py
--- a/jax/experimental/sparse/csr.py
+++ b/jax/experimental/sparse/csr.py
@@ -82,7 +82,6 @@ def _eye(cls, N, M, k, *, dtype=None, index_dtype='int32'):
       # if k is out of range, return an empty matrix.
       return cls._empty((N, M), dtype=dtype, index_dtype=index_dtype)
 
-    k = jnp.asarray(k)
     data = jnp.ones(diag_size, dtype=dtype)
     idx = jnp.arange(diag_size, dtype=index_dtype)
     zero = _const(idx, 0)


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -2493,14 +2493,20 @@ def test_empty(self, cls, shape):
     for k in [-2, 0, 1])
   def test_eye(self, cls, N, M, k):
     sparse_format = cls.__name__.lower()
-    mat = sparse.eye(N, M, k, sparse_format=sparse_format)
+    func = partial(sparse.eye, N, M, k, sparse_format=sparse_format)
     expected = jnp.eye(N, M, k)
     expected_nse = jnp.count_nonzero(expected)
 
+    mat = func()
     self.assertIsInstance(mat, cls)
     self.assertArraysEqual(mat.todense(), expected)
     self.assertEqual(mat.nse, expected_nse)
 
+    mat_jit = jit(func)()
+    self.assertIsInstance(mat_jit, cls)
+    self.assertArraysEqual(mat_jit.todense(), expected)
+    self.assertEqual(mat_jit.nse, expected_nse)
+
   @parameterized.named_parameters(
     {"testcase_name": f"{nse}_BCOO{shape}", "shape": shape, "nse": nse}
     for shape in ([2, 5], [5, 3])