I love JAX … but JAX, frustratingly, doesn’t obey NumPy promotion semantics! Enforcing double precision is understandable from a performance perspective but it breaks NumPy compatibility and frequently introduces subtle and difficult to debug numerical bugs. Nevertheless, JAX has encouraged so much progress in ML software engineering, e.g., I’ve been cribbing from diffrax and torchopt for my PyTorch work.