python3Packages.jax: add operations to cuda test

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!")
   ''