Mathieu Blondel comments

Results 163 comments of


                                            Mathieu Blondel

Getting the stacked outputs of FixedPointIteration.run()

Hi Sophie. Sorry for the late reply. At the moment, we don't have immediate plans to add such a method so it's best that you roll out your own utility...

Implement a `split_rngs` mechanism, like in `flax.linen.scan`

This looks very tricky. Naive question but is the fixed point well-defined when using dropout? Isn't randomness problematic?

document that lower and upper bounds in Bisection need stop_gradient for differentiation

> We could also change the signature of the Bisection method to let lower and upper be taken as parameters I don't think we can. `lower` and `upper` are arguments...

document that lower and upper bounds in Bisection need stop_gradient for differentiation

If you agree with me, we can relabel this issue as documentation. Adding a short paragraph on this would be helpful.

Unnecessary recompilation of _while_loop_lax

CC @froystig

Unnecessary recompilation of _while_loop_lax

what do you propose more concretely?

Remove unused imports

Thanks a lot! Can you merge with main?

Using `vmap` for root finding with a vector of parameters

This comes from the fact that `ScipyMinimize` and `ScipyRootFinding` are not written in pure JAX. We need to wrap the call to `scipy.optimize.minimize` with a `jax.lax.pure_callback` (see #372) as was...

Using `vmap` for root finding with a vector of parameters

@Justin-Tan The modifications are on our side. We need to wrap the call to `scipy.optimize.minimize` or `scipy.optimize.root` with a `pure_callback`.

drop markdown for notebook examples?

Sounds reasonable to me. @jakevdp Any disadvantage?