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