JAX

JAX is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the python3Packages.{jax, jaxlib} packages.

NOTE: JAX requires Python 3.9, the current version of python3 in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).

Example shell.nix, CPU only

let
  pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/eac6215607e4ccceb9631b01ee8f8903a6e82e02.tar.gz")) {};
in pkgs.mkShell {
  buildInputs = with pkgs; [
    python3
    python3Packages.jax
    python3Packages.jaxlib
  ];
}

Example shell.nix with GPU support

JAX defers execution to the jaxlib library for execution. In order to use GPU support you'll need a NVIDIA GPU and OpenGL. In your /etc/nixos/configuration.nix:

# NVIDIA drivers are unfree
nixpkgs.config.allowUnfree = true;
services.xserver.videoDrivers = [ "nvidia" ];
hardware.opengl.enable = true;

Then you can use the jaxlib package by setting the cudaSupport parameter:

let
  # Fork with jax/jaxlib. See https://github.com/NixOS/nixpkgs/pull/134894.
  pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/eac6215607e4ccceb9631b01ee8f8903a6e82e02.tar.gz")) {};
in pkgs.mkShell {
  buildInputs = with pkgs; [
    python3
    python3Packages.jax
    (python3Packages.jaxlib.override { cudaSupport = true; })
  ];
}

You can test that JAX is using the GPU as intended with

python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)"

It should print either "cpu", "gpu", or "tpu".

FAQ

RuntimeError: Unknown: no kernel image is available for execution on the device

This usually indicates that you have a driver version that is too old for the CUDA toolkit version the package is built with. The easiest fix is to set the environment variable XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1". Also consider upgrading your CUDA driver.

See https://github.com/google/jax/issues/5723#issuecomment-913038780.