JAX is a popular deep learning framework that more and more of the NIR-supported libraries are built on.
For PyTorch, we have built the nirtorch package, but no such package exists for JAX.
If you’re interested in developing such a package, please reach out to us!
Either on Discord or by opening an issue.