Non-constant sparse inputs can't be handled because JAX does not allow Scipy sparse matrices as inputs. We could implement a BCOO type explicitly but this would be JAX exclusive, and the user would need to use it from the get go, meaning such graphs would not be compatible with other backends.