From ad4d314c9da708ff810691b6237ab04fba81e4eb Mon Sep 17 00:00:00 2001
From: natsukium <tomoya.otabi@gmail.com>
Date: Sun, 29 Dec 2024 19:18:09 +0900
Subject: [PATCH] python312Packages.jax-cuda12-pjrt: init at 0.4.38

---
 .../jax-cuda12-pjrt/default.nix               | 111 ++++++++++++++++++
 pkgs/top-level/python-packages.nix            |   2 +
 2 files changed, 113 insertions(+)
 create mode 100644 pkgs/development/python-modules/jax-cuda12-pjrt/default.nix

diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
new file mode 100644
index 000000000000..16c0236cc218
--- /dev/null
+++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
@@ -0,0 +1,111 @@
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  fetchurl,
+  autoAddDriverRunpath,
+  autoPatchelfHook,
+  pypaInstallHook,
+  wheelUnpackHook,
+  cudaPackages,
+  python,
+  jaxlib,
+}:
+let
+  inherit (jaxlib) version;
+  inherit (cudaPackages) cudaVersion;
+
+  cudaLibPath = lib.makeLibraryPath (
+    with cudaPackages;
+    [
+      (lib.getLib cuda_cudart) # libcudart.so
+      (lib.getLib cuda_cupti) # libcupti.so
+      (lib.getLib cudnn) # libcudnn.so
+      (lib.getLib libcufft) # libcufft.so
+      (lib.getLib libcusolver) # libcusolver.so
+      (lib.getLib libcusparse) # libcusparse.so
+    ]
+  );
+
+  # Find new releases at https://storage.googleapis.com/jax-releases
+  # 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.
+
+  # upstream does not distribute jax-cuda12-pjrt 0.4.38 binaries for aarch64-linux
+  srcs = {
+    "x86_64-linux" = fetchurl {
+      url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_x86_64.whl";
+      hash = "sha256-g75MWfvPMAd6YAhdmOfVncc4sckeDWKOSsF3n94VrCs=";
+    };
+    "aarch64-linux" = fetchurl {
+      url = "https://storage.googleapis.com/jax-releases/cuda12_plugin/jax_cuda12_pjrt-${version}-py3-none-manylinux2014_aarch64.whl";
+      hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
+    };
+  };
+in
+buildPythonPackage {
+  pname = "jax-cuda12-pjrt";
+  inherit version;
+  pyproject = false;
+
+  src =
+    srcs.${stdenv.hostPlatform.system}
+      or (throw "jax-cuda12-pjrt: No src for ${stdenv.hostPlatform.system}");
+
+  nativeBuildInputs = [
+    autoAddDriverRunpath
+    autoPatchelfHook
+    pypaInstallHook
+    wheelUnpackHook
+  ];
+
+  # The following attributes (buildInputs, postInstall and preInstallCheck) are copied from jaxlib-0.4.28
+  # but it does not recognize GPUs as of 2024-12-29
+  # Dynamic link dependencies
+  buildInputs = [ (lib.getLib stdenv.cc.cc) ];
+
+  # jaxlib looks for ptxas at runtime, eg when running `jax.random.PRNGKey(0)`.
+  # Linking into $out is the least bad solution. See
+  # * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621
+  # * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211
+  # for more info.
+  postInstall = ''
+    mkdir -p $out/${python.sitePackages}/jaxlib/cuda/bin
+    ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jaxlib/cuda/bin/ptxas
+  '';
+
+  # jaxlib contains shared libraries that open other shared libraries via dlopen
+  # and these implicit dependencies are not recognized by ldd or
+  # autoPatchelfHook. That means we need to sneak them into rpath. This step
+  # must be done after autoPatchelfHook and the automatic stripping of
+  # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
+  # patchPhase.
+  preInstallCheck = ''
+    shopt -s globstar
+
+    for file in $out/**/*.so; do
+      echo $file
+      patchelf --add-rpath "${cudaLibPath}" "$file"
+    done
+  '';
+
+  # no tests
+  doCheck = false;
+
+  pythonImportsCheck = [ "jax_plugins.xla_cuda12" ];
+
+  meta = {
+    description = "JAX XLA PJRT Plugin for NVIDIA GPUs";
+    homepage = "https://github.com/jax-ml/jax/tree/main/jax_plugins/cuda";
+    sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ natsukium ];
+    platforms = lib.attrNames srcs;
+    # see CUDA compatibility matrix
+    # https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder
+    broken =
+      !(lib.versionAtLeast cudaVersion "12.1")
+      || !(lib.versionAtLeast cudaPackages.cudnn.version "9.1")
+      || true;
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index e30c9620859c..7d4e039b93bb 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -6560,6 +6560,8 @@ self: super: with self; {
 
   jax = callPackage ../development/python-modules/jax { };
 
+  jax-cuda12-pjrt = callPackage ../development/python-modules/jax-cuda12-pjrt { };
+
   jax-jumpy = callPackage ../development/python-modules/jax-jumpy { };
 
   jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix { };