python3Packages.jax: add operations to cuda test

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

@@ -11,13 +11,21 @@ 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"
 
-    rng = random.PRNGKey(0)
+    rng = random.key(0)
     x = random.normal(rng, (100, 100))
     x @ x
+    jnp.fft.fft(x)
+    jnp.linalg.inv(x)
+
+    y = sparse.CSR.fromdense(x)
+    z = random.normal(rng, (100,))
+    y @ z
 
     print("success!")
   ''