jax

6662da82 - tweaks to simplify masked jaxprs, rnn test