I don't think anyone there had used JAX before? but it looks cool and with FLAX very pythonic, from the quick example I saw,