Rory Conlin

Not sure what lstm.h is, all of the keras2c stuff starts with k2c_, is that a custom wrapper or something?

The `jax.clear_cache()` method seems to work for me with jax v0.4.13, though not quite as well as I'd expect. Using it reduces CPU memory usage by quite a bit but...

If you don't mind a little self promotion, I've been working on a library for quadrature in Jax that should do what you need: https://github.com/f0uriest/quadax And in response to @patrick-kidger...

Thanks! I'm a big fan of equinox so I had considered using your version of `while_loop` but I didn't see it in the documented API so wasn't sure if it...

The original version has 7 possible exit conditions: 1. exceeds maxiter 2. x solve Ax=b to user specified tolerance (ie, no least squares needed, problem appears well posed) 3. x...

Any progress on this? I'd be super interested in having a general abstraction for block linear operators, especially if the blocks themselves can be block linear operators. For example, a...

Can you clarify a bit what you mean by > there isn't (?) a nice way to express a linear-solve-against-a-block-operator in terms of linear solves against the individual operators. Do...

Ok yeah, there would still be a need for specific solvers like this one, but I think the block abstraction would still be useful, and possibly combining it with the...

If box constraints are supported it shouldn't be too much work to add an augmented Lagrangian type method for general constraints. Could maybe be like a "meta solver" where you...

If I remember correctly, IPOPT uses a filter to decide on accepting steps, which might be difficult to implement in JAX with compile time memory bounds? It can also switch...