Dxyk
commented
2 years ago Update - I was able to get it to work by jitting the CG solver and marking the A and M operators as static.
This is inspired by this discussion https://github.com/google/jax/discussions/9804, but I don't quite understand why marking the functions as static would resolve the TypeError of Zero(ShapedArray())
@jit()
def step():
# Other operations that involve x and b ...
def A_operator(x: jnp.ndarray) -> jnp.ndarray:
# compute res using convolution. Left unchanged
return res
def M_operator(b: jnp.ndarray) -> jnp.ndarray:
# Set as identity matrix
return b
# original version that fails grad with TypeError
#jsp.sparse.linalg.cg(A=A_operator, b=b, M=M_operator))
# new version that works with grad and jit
jax.jit(jsp.sparse.linalg.cg, static_argnames=('A', 'M'))(A=A_operator, b=b, M=M_operator)
zhangqiaorjc
Hello, I have been using Jax to build a differentiable fluid simulator, which involves solving a sparse linear system using the Conjugate Gradient algorithm (
jax.scipy.sparse.linalg.cg). When I tried to take gradient against the jitted simulator (grad(jit(step)), wherestepis one step of the simulation that contains one CG solve), I encountered the following TypeError.Interestingly, without jitting (
grad(step)only) and keeping everything else the same, the TypeError will not be thrown, and the gradient is computed correctly.When searching for similar issues, I found #1536 with the same error and might be related.
Unfortunately, I was unable to come up with a minimal example that can reproduce this error. If it is necessary, I'd be glad to provide access to my repo.
Here are my package versions
jax==0.3.13,jaxlib==0.3.10.Please see the stack trace attached below. Thanks for your help in advance!
Stack Trace
```text Traceback (most recent call last): File ".../FluidSim/demos/simulation_grad_bug_demo.py", line 61, in