python3Packages.jax: add operations to cuda test

Author: stephen-huan

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

@@ -95,7 +95,7 @@ buildPythonPackage {
     wheelUnpackHook
   ];
 
-  # jax-cuda12-plugin looks for ptxas at runtime, e.g. with a xla custom call.
+  # jax-cuda12-plugin looks for ptxas at runtime, e.g. with a triton kernel.
   # Linking into $out is the least bad solution. See
   # * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621
   # * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211

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

@@ -11,13 +11,18 @@ pkgs.writers.writePython3Bin "jax-test-cuda"
   }
   ''
     import jax
+    import jax.numpy as jnp
     from jax import random
+    from jax.experimental import sparse
 
-    assert jax.devices()[0].platform == "gpu"
+    assert jax.devices()[0].platform == "gpu"  # libcuda.so
 
-    rng = random.PRNGKey(0)
+    rng = random.key(0)  # libcudart.so, libcudnn.so
     x = random.normal(rng, (100, 100))
-    x @ x
+    x @ x  # libcublas.so
+    jnp.fft.fft(x)  # libcufft.so
+    jnp.linalg.inv(x)  # libcusolver.so
+    sparse.CSR.fromdense(x) @ x  # libcusparse.so
 
     print("success!")
   ''