30 lines
594 B
Nix
30 lines
594 B
Nix
{
|
|
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!")
|
|
''
|