pyrox.dev
1#! /usr/bin/env nix-shell
2#! nix-shell -i sh -p jq
3
4prefetch() {
5 expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = true; }).python$1.pkgs.$3.src.url"
6 url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r)
7 echo "$url"
8 sha256=$(nix-prefetch-url "$url")
9 nix --extra-experimental-features nix-command hash convert --to sri --hash-algo sha256 "$sha256"
10 echo
11}
12
13for py in "311" "312" "313"; do
14 prefetch "$py" "x86_64-linux" "jaxlib-bin"
15 prefetch "$py" "aarch64-linux" "jaxlib-bin"
16 prefetch "$py" "aarch64-darwin" "jaxlib-bin"
17 prefetch "$py" "x86_64-linux" "jax-cuda12-plugin"
18 prefetch "$py" "aarch64-linux" "jax-cuda12-plugin"
19done
20
21for arch in "x86_64-linux" "aarch64-linux"; do
22 prefetch "312" "$arch" "jax-cuda12-pjrt"
23done