jax: init at 0.2.19, jaxlib: init at 0.1.71

[samuela]

Motivation for this change

Packaging of the JAX (https://github.com/google/jax) library. JAX is a modern framework for differentiable programming. The jax package uses jaxlib as it's "backend". Currently only the CPU backend running on x86_64-linux is packaged.

TODO, unfinished items left to future PRs:

• Fix jnp.array([1, 2, 3]).
• Run the jax test suite.
• GPU backend support
• TPU backend support
• Fix the build on macOS. Currently segfaults when running the test suite.

Useful resources:

• https://jax.readthedocs.io/en/latest/developer.html#running-the-tests
• https://github.com/NixOS/nixpkgs/blob/master/doc/languages-frameworks/python.section.md

Things done

Created derivations for the jax and jaxlib python packages.

• Tested using sandboxing (nix.useSandbox on NixOS, or option sandbox in nix.conf on non-NixOS linux)
• Built on platform(s)
  • NixOS
  • macOS
  • other Linux distributions
• Tested via one or more NixOS test(s) if existing and applicable for the change (look inside nixos/tests)
• Tested compilation of all packages that depend on this change using nix-shell -p nixpkgs-review --run "nixpkgs-review wip"
• Tested execution of all binary files (usually in ./result/bin/)
• 21.11 Release Notes (or backporting 21.05 Release notes)
  • (Package updates) Added a release notes entry if the change is major or breaking
  • (Module updates) Added a release notes entry if the change is significant
  • (Module addition) Added a release notes entry if adding a new NixOS module
• Fits CONTRIBUTING.md.

[nixos-discourse]

This pull request has been mentioned on NixOS Discourse. There might be relevant details there:

https://discourse.nixos.org/t/help-packaging-jax/14613/1

[samuela]

On my (outdated) linux machine I'm getting:

>>> import jax
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/nix/store/994kvxz1n7qch3zn8gfxi9wdd7nkzlm7-python3.9-jax-0.2.19/lib/python3.9/site-packages/jax/__init__.py", line 37, in <module>
    from . import config as _config_module
  File "/nix/store/994kvxz1n7qch3zn8gfxi9wdd7nkzlm7-python3.9-jax-0.2.19/lib/python3.9/site-packages/jax/config.py", line 18, in <module>
    from jax._src.config import config
  File "/nix/store/994kvxz1n7qch3zn8gfxi9wdd7nkzlm7-python3.9-jax-0.2.19/lib/python3.9/site-packages/jax/_src/config.py", line 27, in <module>
    from jax import lib
  File "/nix/store/994kvxz1n7qch3zn8gfxi9wdd7nkzlm7-python3.9-jax-0.2.19/lib/python3.9/site-packages/jax/lib/__init__.py", line 71, in <module>
    cpu_feature_guard.check_cpu_features()
RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source.

So it may be worth considering building jaxlib from source but honestly that also sounds like a PITA.

[samuela]

Attempted to use autoPatchelfHook in the building of jaxlib in 59d6dd2. Not sure that it's working however: https://discourse.nixos.org/t/autopatchelfhook-not-patching-all-dependencies/14634/2.

@SuperSandro2000 Thanks for reviewing! I'll start incorporating your suggestions once I can get it to build and pass tests!

[nixos-discourse]

This pull request has been mentioned on NixOS Discourse. There might be relevant details there:

https://discourse.nixos.org/t/prs-ready-for-review/3032/582

[SuperSandro2000]

Those comments are not required and they can get outdated quite fast.

[nixos-discourse]

This pull request has been mentioned on NixOS Discourse. There might be relevant details there:

https://discourse.nixos.org/t/prs-ready-for-review/3032/584

[SuperSandro2000]

Result of nixpkgs-review pr 134894 run on x86_64-linux 1

2 packages built:

• python39Packages.jax
• python39Packages.jaxlib