colin/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.jax: add CUDA tests in passthru

Browse Source 
Organize CUDA-enabled tests for jax and jaxlib into passthru scripts to facilitate testing.
This commit is contained in:
Samuel Ainsworth 2024-02-28 02:50:45 +00:00
parent 9e8cdafeb8
commit 2f4ac6f803
3 changed files with 36 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
pkgs/development/python-modules/jax/default.nix
View File

@ -1,6 +1,7 @@
{ 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";

17
pkgs/development/python-modules/jax/test-cuda.nix Normal file
View File

@ -0,0 +1,17 @@
{ 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!")
''

11
pkgs/development/python-modules/jaxlib/bin.nix
View File

@ -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