Using sparse solver within jitted function

[fakemade]

Hello JAX team and users. I have a question about sparse solver. I have the follow jitted code (this is example of real code):

K_d = sparse.BCOO.fromdense(K)
u_free = jax.scipy.sparse.linalg.gmres(K_d,self.FVector[self.freeDofs])[0]
u = jnp.zeros((self.TypeTask.value * self.NodeCount))
u = u.at[self.freeDofs].add(u_free.reshape(-1))  # homog bc wherev fixed
return u

When i run this code, i catch error about 'nse' argument, that it's must be predefined. But i also can't use this row

import jax.numpy as jnp
cnt = jnp.count_nonzero(K)
K_d = sparse.BCOO.fromdense(K, nse=cnt)

It is also give me error, that somethings are depended from 'K'. How can i use sparse solver wthin jit function?

P.S. Sorry, i am juniour in python programming and i don`t know, how to ask a question correctly. If you need more information i am ready to get it to you.

[jakevdp]

The issue is that you need a static value of nse, in order for all your arrays to be statically shaped (see JAX Sharp Bits: Dynamic Shapes). You can do this by defining nse outside the context of jit, vmap, or another JAX transformation.

It's hard to be much more specific than that without a reproduction of the code you're running – the partial snippets you included don't really provide enough information to give more specific advice. You might think about adding a minimal reproducible example to your question.