flax

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