python3Packages.jaxlib-bin: support Python 3.9-3.12

  # official instructions recommend installing CPU-only versions via PyPI.

  cpuSrcs =

    let

      getSrcFromPypi = { platform, hash }: fetchPypi {

        inherit version platform hash;

      getSrcFromPypi = { platform, dist, hash }: fetchPypi {

        inherit version platform dist hash;

        pname = "jaxlib";

        format = "wheel";

        # See the `disabled` attr comment below.

        dist = "cp310";

        python = "cp310";

        abi = "cp310";

        python = dist;

        abi = dist;

      };

    in

    {

      "x86_64-linux" = getSrcFromPypi {

      "3.9-x86_64-linux" = getSrcFromPypi {

        platform = "manylinux2014_x86_64";

        dist = "cp39";

        hash = "sha256-8bTrWutuK0qVnbkcwMfgBf414YdaLc3GK5IsCm/JNPE=";

      };

      "3.9-aarch64-darwin" = getSrcFromPypi {

        platform = "macosx_11_0_arm64";

        dist = "cp39";

        hash = "sha256-Tmv2iOqlNbZqw/rYjef6GmM0N18EA5JTt6T3lQe+4Rs=";

      };

      "3.9-x86_64-darwin" = getSrcFromPypi {

        platform = "macosx_10_14_x86_64";

        dist = "cp39";

        hash = "sha256-mDT1INLqPdCkxtMMFR0qHLOIZdWEy8Iuzw1/vOoECsA=";

      };

      "3.10-x86_64-linux" = getSrcFromPypi {

        platform = "manylinux2014_x86_64";

        dist = "cp310";

        hash = "sha256-ksnY+CPEstact5lKjbSg+ZSPJtSt0Y0NFWEFufBCByk=";

      };

      "aarch64-darwin" = getSrcFromPypi {

      "3.10-aarch64-darwin" = getSrcFromPypi {

        platform = "macosx_11_0_arm64";

        dist = "cp310";

        hash = "sha256-O7dHvdKLKfNELGfF4TKy7N5EX6Ca7Zu8OtLXWvFykR8=";

      };

      "x86_64-darwin" = getSrcFromPypi {

      "3.10-x86_64-darwin" = getSrcFromPypi {

        platform = "macosx_10_14_x86_64";

        dist = "cp310";

        hash = "sha256-gqKMUZSXrt8sQtTAoQbzAfCzO8gM9Y1/tZpuJVWyN0Y=";

      };

      "3.11-x86_64-linux" = getSrcFromPypi {

        platform = "manylinux2014_x86_64";

        dist = "cp311";

        hash = "sha256-m+NDzwXMNboNjDl2nLY+vqAoN2dQJZVWb1UQDpqqDPw=";

      };

      "3.11-aarch64-darwin" = getSrcFromPypi {

        platform = "macosx_11_0_arm64";

        dist = "cp311";

        hash = "sha256-zCOAjaWWCQT9Jnm1jjc1Rh5gemqy7ACtTKLM0MqSJzM=";

      };

      "3.11-x86_64-darwin" = getSrcFromPypi {

        platform = "macosx_10_14_x86_64";

        dist = "cp311";

        hash = "sha256-gOLIxkk+2hew2GqWu1WgMVEx1YEutx7Zod7QbwsuUVQ=";

      };

      "3.12-x86_64-linux" = getSrcFromPypi {

        platform = "manylinux2014_x86_64";

        dist = "cp312";

        hash = "sha256-BZTmkgNuV4nWtfbY4t/19aP43szZQEdgpFXh5qwGRXk=";

      };

      "3.12-aarch64-darwin" = getSrcFromPypi {

        platform = "macosx_11_0_arm64";

        dist = "cp312";

        hash = "sha256-aAMTrLXU9EYwPv+kdeyI88/D7b4NANB39Fn8vuXUqFA=";

      };

      "3.12-x86_64-darwin" = getSrcFromPypi {

        platform = "macosx_10_14_x86_64";

        dist = "cp312";

        hash = "sha256-KHzlIfa9KtYcHX+i/F/SKaYTpD4/XjHVu5j3BdRTUmc=";

      };

    };

  # Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html.

  # When upgrading, you can get these hashes from prefetch.sh. See

  # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index.

  gpuSrc = fetchurl {

  gpuSrcs = {

    "3.9" = fetchurl {

      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl";

      hash = "sha256-WB5Vbr/XeYKXCP/3DIXF20jR6/1xE3huX1h5ow8ETl0=";

    };

    "3.10" = fetchurl {

      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl";

      hash = "sha256-zfN0n31+5GohwBkeQrqHus4qOyhM/GEdqG6KUupCZ4o=";

    };

    "3.11" = fetchurl {

      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl";

      hash = "sha256-Q8ZtF2GCrG30GFbCeCZTWPmW2TBybeXzh2u+NRiYpx4=";

    };

    "3.12" = fetchurl {

      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl";

      hash = "sha256-lphkSDOJ9SwbO0hp/xC1bYn5fWgth9A9Iwsc9zV0buI=";

    };

  };

in

buildPythonPackage {

  inherit version;

  format = "wheel";

  disabled = !(pythonVersion == "3.10");

  disabled = !(pythonVersion == "3.9" || pythonVersion == "3.10" || pythonVersion == "3.11" || pythonVersion == "3.12");

  # See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.

  src =

    if !cudaSupport then

      (

        cpuSrcs."${stdenv.hostPlatform.system}"

        cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}"

          or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")

      ) else gpuSrc;

      ) else gpuSrcs."${pythonVersion}";

  # Prebuilt wheels are dynamically linked against things that nix can't find.

  # Run `autoPatchelfHook` to automagically fix them.

#!/usr/bin/env bash

#! /usr/bin/env nix-shell

#! nix-shell -i sh -p jq

prefetch () {

    expr="(import <nixpkgs> { system = \"$1\"; config.cudaSupport = $2; }).python3.pkgs.jaxlib-bin.src.url"

    expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = $3; }).python$1.pkgs.jaxlib-bin.src.url"

    url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r)

    echo "$url"

    sha256=$(nix-prefetch-url "$url")

    echo

}

prefetch "x86_64-linux" "false"

prefetch "aarch64-darwin" "false"

prefetch "x86_64-darwin" "false"

prefetch "x86_64-linux" "true"

for py in "39" "310" "311" "312"; do

    prefetch "$py" "x86_64-linux" "false"

    prefetch "$py" "aarch64-darwin" "false"

    prefetch "$py" "x86_64-darwin" "false"

    prefetch "$py" "x86_64-linux" "true"

done