JAX is the new deep learning kid on the block. I use it in my day-to-day work and I’ve always wanted to get it packaged nicely on Nix. I took a stab at it here: jax: init at 0.2.19, jaxlib: init at 0.1.71 by samuela · Pull Request #134894 · NixOS/nixpkgs · GitHub.
But unfortunately it’s full of segfaults at the moment. Any help or advice would be greatly appreciated!