I am actually not surprised by JAX converting better to XLA. Also deep respect for anybody in this space as their is lot of complexity (?) to deal with at the framework and compiler level.