From 2f4ac6f80307f4f66bd7491bc2cd9e690e5ee81d Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Wed, 28 Feb 2024 02:50:45 +0000
Subject: [PATCH] python3Packages.jax: add CUDA tests in passthru

Organize CUDA-enabled tests for jax and jaxlib into passthru scripts to facilitate testing.
---
 .../development/python-modules/jax/default.nix | 18 ++++++++++++++++++
 .../python-modules/jax/test-cuda.nix           | 17 +++++++++++++++++
 pkgs/development/python-modules/jaxlib/bin.nix | 11 +----------
 3 files changed, 36 insertions(+), 10 deletions(-)
 create mode 100644 pkgs/development/python-modules/jax/test-cuda.nix

diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
index 0450f73e639b..d41987c9a539 100644
--- a/pkgs/development/python-modules/jax/default.nix
+++ b/pkgs/development/python-modules/jax/default.nix
@@ -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";
diff --git a/pkgs/development/python-modules/jax/test-cuda.nix b/pkgs/development/python-modules/jax/test-cuda.nix
new file mode 100644
index 000000000000..d156061f3849
--- /dev/null
+++ b/pkgs/development/python-modules/jax/test-cuda.nix
@@ -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!")
+''
diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix
index 5b6fff2d7c0e..199b352b2aa2 100644
--- a/pkgs/development/python-modules/jaxlib/bin.nix
+++ b/pkgs/development/python-modules/jaxlib/bin.nix
@@ -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