flax

472b82a2 - Improve partial_eval_by_shape compatibility with https://github.com/google/jax/pull/3370 and custom derivative rules.