flax

4d057dd1 - Avoid depending on JAX internals, which are about to change (see cl/681582933).