Something went wrong. Try again.
lol
Something went wrong. Try again.
594 B · 29 lines
Nix
at master
123456789101112131415161718192021222324252627282930{ jax, pkgs,}:
pkgs.writers.writePython3Bin "jax-test-cuda" { libraries = [ jax ] ++ jax.optional-dependencies.cuda; } '' import jax import jax.numpy as jnp from jax import random from jax.experimental import sparse
assert jax.devices()[0].platform == "gpu" # libcuda.so
rng = random.key(0) # libcudart.so, libcudnn.so x = random.normal(rng, (100, 100)) 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!") ''