natsukium
1802b4b4e5
Diff: https://github.com/jax-ml/jax/compare/jax-v0.4.28...jax-v0.4.38 Changelog: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.38
24 lines
343 B
Nix
24 lines
343 B
Nix
{
|
|
jax,
|
|
pkgs,
|
|
}:
|
|
|
|
pkgs.writers.writePython3Bin "jax-test-cuda"
|
|
{
|
|
libraries = [
|
|
jax
|
|
] ++ jax.optional-dependencies.cuda;
|
|
}
|
|
''
|
|
import jax
|
|
from jax import random
|
|
|
|
assert jax.devices()[0].platform == "gpu"
|
|
|
|
rng = random.PRNGKey(0)
|
|
x = random.normal(rng, (100, 100))
|
|
x @ x
|
|
|
|
print("success!")
|
|
''
|