3 ms·

IIRC jax's `scipy.sparse.linalg.bicgstab` does support multiple right hand sides. EDIT: Or rather, all the solvers under jax's `scipy.sparse.linalg` all suppor

• bee_rider 2y ago
  Oh dang, that’s pretty awesome, thanks. “array or tree of arrays” sounds very general, probably even better than an old fashioned 2D array.
  • trostaft 2y ago
    'tree of arrays' Ahh, that's just Jax's concept of pytrees. It was something that they invented to make it easier (this is how I view it, not complete) to pass complex objects to function but still be able to easily consider them as a concatenated vector for AD etc.. E.g. a common pattern is to pass parameters `p` to a function and then internally break them into their physical interpretations, e.g. `mass = p[0]`, `velocity = p[1]`. Pytrees let you just use something like a dictionary `p = {'mass' = 1.0, 'velocity = 1.0'}`, which is a stylistically more natural structure to pass around, and then jax is structured to understand later when AD'ing or otherwise that you're doing so with respect to the 'leaves' of the tree, or the values of the mass and velocity. Hopefully someone corrects me if I'm not right about this. I'm hardly 100% on Jax's vision on PyTrees. As an aside, just a list of right hand sides `[b1, b2, ..., bm]` is valid.