flax

Ignore `jax.host_{id,count}()` deprecation warning

#1301

Merged