jax

c04dea1c - Begin implementing mask(jit)