Developing JAX extensions

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.