flax

81efeef0 - Remove outdated compatibility code.

Remove outdated compatibility code. The minimum supported JAX version is now 0.4.27, so these compatibility shims for older versions are no longer necessary.