Commit 2f4ac6f8 authored by Samuel Ainsworth's avatar Samuel Ainsworth
Browse files

python3Packages.jax: add CUDA tests in passthru 

Organize CUDA-enabled tests for jax and jaxlib into passthru scripts to facilitate testing.
parent 9e8cdafe

pkgs/development/python-modules/jax/default.nix

+18 −0
Original line number Diff line number Diff line 
{ lib

, blas

, buildPythonPackage

, callPackage

, setuptools

, importlib-metadata

, fetchFromGitHub
@@ -126,6 +127,23 @@ buildPythonPackage rec { 


  pythonImportsCheck = [ "jax" ];



  # Test CUDA-enabled jax and jaxlib. Running CUDA-enabled tests is not

  # currently feasible within the nix build environment so we have to maintain

  # this script separately. See https://github.com/NixOS/nixpkgs/pull/256230

  # for a possible remedy to this situation.

  #

  # Run these tests with eg

  #

  #   NIXPKGS_ALLOW_UNFREE=1 nixglhost -- nix run --impure .#python3Packages.jax.passthru.tests.test_cuda_jaxlibBin

  passthru.tests = {

    test_cuda_jaxlibSource = callPackage ./test-cuda.nix {

      jaxlib = jaxlib.override { cudaSupport = true; };

    };

    test_cuda_jaxlibBin = callPackage ./test-cuda.nix {

      jaxlib = jaxlib-bin.override { cudaSupport = true; };

    };

  };



  meta = with lib; {

    description = "Differentiate, compile, and transform Numpy code";

    homepage = "https://github.com/google/jax";

pkgs/development/python-modules/jax/test-cuda.nix

0 → 100644
+17 −0
Original line number Diff line number Diff line 
{ jax

, jaxlib

, pkgs

}:



pkgs.writers.writePython3Bin "jax-test-cuda" { libraries = [ jax jaxlib ]; } ''

  import jax

  from jax import random



  assert jax.devices()[0].platform == "gpu"



  rng = random.PRNGKey(0)

  x = random.normal(rng, (100, 100))

  x @ x



  print("success!")

''

pkgs/development/python-modules/jaxlib/bin.nix

+1 −10
Original line number Diff line number Diff line
@@ -2,16 +2,7 @@ 
# backend will require some additional work. Those wheels are located here:

# https://storage.googleapis.com/jax-releases/libtpu_releases.html.



# For future reference, the easiest way to test the GPU backend is to run

#   NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib-bin.override { cudaSupport = true; }"

#   export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1

#   python -c "from jax.lib import xla_bridge; assert xla_bridge.get_backend().platform == 'gpu'"

#   python -c "from jax import random; random.PRNGKey(0)"

#   python -c "from jax import random; x = random.normal(random.PRNGKey(0), (100, 100)); x @ x"

# There's no convenient way to test the GPU backend in the derivation since the

# nix build environment blocks access to the GPU. See also:

#   * https://github.com/google/jax/issues/971#issuecomment-508216439

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

# See `python3Packages.jax.passthru` for CUDA tests.



{ absl-py

, autoPatchelfHook