mirror of
https://github.com/NixOS/nixpkgs.git
synced 2025-01-26 23:03:41 +00:00
2f4ac6f803
Organize CUDA-enabled tests for jax and jaxlib into passthru scripts to facilitate testing.
18 lines
285 B
Nix
18 lines
285 B
Nix
{ jax
|
|
, jaxlib
|
|
, pkgs
|
|
}:
|
|
|
|
pkgs.writers.writePython3Bin "jax-test-cuda" { libraries = [ jax jaxlib ]; } ''
|
|
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!")
|
|
''
|