diff --git a/nixos/doc/manual/release-notes/rl-2505.section.md b/nixos/doc/manual/release-notes/rl-2505.section.md
index 6f1453de55a4..15497d928149 100644
--- a/nixos/doc/manual/release-notes/rl-2505.section.md
+++ b/nixos/doc/manual/release-notes/rl-2505.section.md
@@ -258,6 +258,11 @@
 
 - `siduck76-st` has been renamed to `st-snazzy`, like the project's [flake](https://github.com/siduck/st/blob/main/flake.nix).
 
+- `python3Packages.jax` now directly depends on `python3Packages.jaxlib`.
+  As a result, packages that depend on jax no longer need to include jaxlib to their dependencies.
+  There is also a breaking change in the handling of CUDA. Instead of using a CUDA compatible jaxlib
+  as before, you can use plugins like `python3Packages.jax-cuda12-plugin`.
+
 
 <!-- To avoid merge conflicts, consider adding your item at an arbitrary place in the list instead. -->
 
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/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
new file mode 100644
index 000000000000..d672a1c82e1a
--- /dev/null
+++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
@@ -0,0 +1,117 @@
+{
+  lib,
+  stdenv,
+  buildPythonPackage,
+  fetchPypi,
+  autoAddDriverRunpath,
+  autoPatchelfHook,
+  pypaInstallHook,
+  wheelUnpackHook,
+  cudaPackages,
+  python,
+  jaxlib,
+  jax-cuda12-pjrt,
+}:
+let
+  inherit (cudaPackages) cudaVersion;
+  inherit (jaxlib) version;
+
+  getSrcFromPypi =
+    {
+      platform,
+      dist,
+      hash,
+    }:
+    fetchPypi {
+      inherit
+        version
+        platform
+        dist
+        hash
+        ;
+      pname = "jax_cuda12_plugin";
+      format = "wheel";
+      python = dist;
+      abi = dist;
+    };
+
+  # upstream does not distribute jax-cuda12-plugin 0.4.38 binaries for aarch64-linux
+  srcs = {
+    "3.10-x86_64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_x86_64";
+      dist = "cp310";
+      hash = "sha256-nULpmc1k3VZ8FJ7Wj3k5K6iGRDZCGLtjbNzvoBl8kv4=";
+    };
+    "3.10-aarch64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_aarch64";
+      dist = "cp310";
+      hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
+    };
+    "3.11-x86_64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_x86_64";
+      dist = "cp311";
+      hash = "sha256-cEZUOG8OYAoCgdquqViCqmekfttoOTthsbFzx+jKdKg=";
+    };
+    "3.11-aarch64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_aarch64";
+      dist = "cp311";
+      hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
+    };
+    "3.12-x86_64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_x86_64";
+      dist = "cp312";
+      hash = "sha256-Ufas/3Ew63LrsCU039NYGg9eoGlx3lLX68Ia1Nh/5x4=";
+    };
+    "3.12-aarch64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_aarch64";
+      dist = "cp312";
+      hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
+    };
+    "3.13-x86_64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_x86_64";
+      dist = "cp313";
+      hash = "sha256-CSKKTCtEO3aozZqOwikGAInEzINuBiSWh1ptb9xm0x8=";
+    };
+    "3.13-aarch64-linux" = getSrcFromPypi {
+      platform = "manylinux2014_aarch64";
+      dist = "cp313";
+      hash = "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=";
+    };
+  };
+in
+buildPythonPackage {
+  pname = "jax-cuda12-plugin";
+  inherit version;
+  pyproject = false;
+
+  src = (
+    srcs."${python.pythonVersion}-${stdenv.hostPlatform.system}"
+      or (throw "python${python.pythonVersion}Packages.jax-cuda12-plugin is not supported on ${stdenv.hostPlatform.system}")
+  );
+
+  nativeBuildInputs = [
+    autoAddDriverRunpath
+    autoPatchelfHook
+    pypaInstallHook
+    wheelUnpackHook
+  ];
+
+  dependencies = [ jax-cuda12-pjrt ];
+
+  pythonImportsCheck = [ "jax_cuda12_plugin" ];
+
+  meta = {
+    description = "JAX Plugin for CUDA12";
+    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.platforms.linux;
+    # 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/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix
index ba9a621b1eab..99154a8ea86a 100644
--- a/pkgs/development/python-modules/jax/default.nix
+++ b/pkgs/development/python-modules/jax/default.nix
@@ -1,25 +1,38 @@
 {
   lib,
+  config,
+  stdenv,
   blas,
-  buildPythonPackage,
-  callPackage,
-  setuptools,
-  importlib-metadata,
-  fetchFromGitHub,
-  jaxlib,
-  jaxlib-bin,
-  jaxlib-build,
-  hypothesis,
   lapack,
-  matplotlib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  cudaSupport ? config.cudaSupport,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  jaxlib,
   ml-dtypes,
   numpy,
   opt-einsum,
+  scipy,
+
+  # optional-dependencies
+  jax-cuda12-plugin,
+
+  # tests
+  cloudpickle,
+  hypothesis,
+  matplotlib,
   pytestCheckHook,
   pytest-xdist,
-  pythonOlder,
-  scipy,
-  stdenv,
+
+  # passthru
+  callPackage,
+  jax,
+  jaxlib-build,
+  jaxlib-bin,
 }:
 
 let
@@ -27,38 +40,41 @@ let
 in
 buildPythonPackage rec {
   pname = "jax";
-  version = "0.4.28";
+  version = "0.4.38";
   pyproject = true;
 
-  disabled = pythonOlder "3.9";
-
   src = fetchFromGitHub {
     owner = "google";
     repo = "jax";
     # google/jax contains tags for jax and jaxlib. Only use jax tags!
-    rev = "refs/tags/jax-v${version}";
-    hash = "sha256-qSHPwi3is6Ts7pz5s4KzQHBMbcjGp+vAOsejW3o36Ek=";
+    tag = "jax-v${version}";
+    hash = "sha256-H8I9Mkz6Ia1RxJmnuJOSevLGHN2J8ey59ZTlFx8YfnA=";
   };
 
-  nativeBuildInputs = [ setuptools ];
+  build-system = [ setuptools ];
 
   # The version is automatically set to ".dev" if this variable is not set.
   # https://github.com/google/jax/commit/e01f2617b85c5bdffc5ffb60b3d8d8ca9519a1f3
   JAX_RELEASE = "1";
 
-  # jaxlib is _not_ included in propagatedBuildInputs because there are
-  # different versions of jaxlib depending on the desired target hardware. The
-  # JAX project ships separate wheels for CPU, GPU, and TPU.
-  propagatedBuildInputs = [
+  dependencies = [
+    jaxlib
     ml-dtypes
     numpy
     opt-einsum
     scipy
-  ] ++ lib.optional (pythonOlder "3.10") importlib-metadata;
+  ] ++ lib.optionals cudaSupport optional-dependencies.cuda;
+
+  optional-dependencies = rec {
+    cuda = [ jax-cuda12-plugin ];
+    cuda12 = cuda;
+    cuda12_pip = cuda;
+    cuda12_local = cuda;
+  };
 
   nativeCheckInputs = [
+    cloudpickle
     hypothesis
-    jaxlib
     matplotlib
     pytestCheckHook
     pytest-xdist
@@ -71,11 +87,17 @@ buildPythonPackage rec {
   # which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2.
   # Not a big deal, this is how the JAX docs suggest running the test suite
   # anyhow.
-  pytestFlagsArray = [
-    "--numprocesses=4"
-    "-W ignore::DeprecationWarning"
-    "tests/"
-  ];
+  pytestFlagsArray =
+    [
+      "--numprocesses=4"
+      "-W ignore::DeprecationWarning"
+      "tests/"
+    ]
+    ++ lib.optionals stdenv.hostPlatform.isDarwin [
+      # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
+      "--deselect tests/shape_poly_test.py::ShapePolyTest"
+      "--deselect tests/tree_util_test.py::TreeTest"
+    ];
 
   # Prevents `tests/export_back_compat_test.py::CompatTest::test_*` tests from failing on darwin with
   # PermissionError: [Errno 13] Permission denied: '/tmp/back_compat_testdata/test_*.py'
@@ -125,16 +147,27 @@ buildPythonPackage rec {
       # Fails on some hardware due to some numerical error
       # See https://github.com/google/jax/issues/18535
       "testQdwhWithOnRankDeficientInput5"
+    ]
+    ++ lib.optionals stdenv.hostPlatform.isDarwin [
+      # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
+      "testInAxesPyTreePrefixMismatchError"
+      "testInAxesPyTreePrefixMismatchErrorKwargs"
+      "testOutAxesPyTreePrefixMismatchError"
+      "test_tree_map"
+      "test_vjp_rule_inconsistent_pytree_structures_error"
+      "test_vmap_in_axes_tree_prefix_error"
+      "test_vmap_mismatched_axis_sizes_error_message_issue_705"
     ];
 
-  disabledTestPaths = [
-    # Segmentation fault. See https://gist.github.com/zimbatm/e9b61891f3bcf5e4aaefd13f94344fba
-    "tests/linalg_test.py"
-  ]
-  ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
-    # RuntimeWarning: invalid value encountered in cast
-    "tests/lax_test.py"
-  ];
+  disabledTestPaths =
+    [
+      # Segmentation fault. See https://gist.github.com/zimbatm/e9b61891f3bcf5e4aaefd13f94344fba
+      "tests/linalg_test.py"
+    ]
+    ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isAarch64) [
+      # RuntimeWarning: invalid value encountered in cast
+      "tests/lax_test.py"
+    ];
 
   pythonImportsCheck = [ "jax" ];
 
@@ -147,25 +180,26 @@ buildPythonPackage rec {
   #
   #   NIXPKGS_ALLOW_UNFREE=1 nixglhost -- nix run --impure .#python3Packages.jax.passthru.tests.test_cuda_jaxlibBin
   passthru.tests = {
-    test_cuda_jaxlibSource = callPackage ./test-cuda.nix {
-      jaxlib = jaxlib-build.override { cudaSupport = true; };
-    };
+    # jaxlib-build is broken as of 2024-12-20
+    # test_cuda_jaxlibSource = callPackage ./test-cuda.nix {
+    #   jax = jax.override { jaxlib = jaxlib-build; };
+    # };
     test_cuda_jaxlibBin = callPackage ./test-cuda.nix {
-      jaxlib = jaxlib-bin.override { cudaSupport = true; };
+      jax = jax.override { jaxlib = jaxlib-bin; };
     };
   };
 
   # updater fails to pick the correct branch
   passthru.skipBulkUpdate = true;
 
-  meta = with lib; {
+  meta = {
     description = "Source-built JAX frontend: differentiate, compile, and transform Numpy code";
     longDescription = ''
       This is the JAX frontend package, it's meant to be used together with one of the jaxlib implementations,
       e.g. `python3Packages.jaxlib`, `python3Packages.jaxlib-bin`, or `python3Packages.jaxlibWithCuda`.
     '';
     homepage = "https://github.com/google/jax";
-    license = licenses.asl20;
-    maintainers = with maintainers; [ samuela ];
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ samuela ];
   };
 }
diff --git a/pkgs/development/python-modules/jax/test-cuda.nix b/pkgs/development/python-modules/jax/test-cuda.nix
index 5aca523f3177..ebeebb7ed88c 100644
--- a/pkgs/development/python-modules/jax/test-cuda.nix
+++ b/pkgs/development/python-modules/jax/test-cuda.nix
@@ -1,6 +1,5 @@
 {
   jax,
-  jaxlib,
   pkgs,
 }:
 
@@ -8,8 +7,7 @@ pkgs.writers.writePython3Bin "jax-test-cuda"
   {
     libraries = [
       jax
-      jaxlib
-    ];
+    ] ++ jax.optional-dependencies.cuda;
   }
   ''
     import jax
diff --git a/pkgs/development/python-modules/jaxlib/bin.nix b/pkgs/development/python-modules/jaxlib/bin.nix
index 5b7fb2d88a93..dd82358f03b7 100644
--- a/pkgs/development/python-modules/jaxlib/bin.nix
+++ b/pkgs/development/python-modules/jaxlib/bin.nix
@@ -6,46 +6,24 @@
 
 {
   absl-py,
-  autoAddDriverRunpath,
   autoPatchelfHook,
   buildPythonPackage,
-  config,
   fetchPypi,
-  fetchurl,
   flatbuffers,
-  jaxlib-build,
   lib,
   ml-dtypes,
   python,
   scipy,
   stdenv,
-  # Options:
-  cudaSupport ? config.cudaSupport,
-  cudaPackages,
 }:
 
 let
-  inherit (cudaPackages) cudaVersion;
-
-  version = "0.4.28";
-
+  version = "0.4.38";
   inherit (python) pythonVersion;
 
-  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
-    ]
-  );
-
   # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the
   # official instructions recommend installing CPU-only versions via PyPI.
-  cpuSrcs =
+  srcs =
     let
       getSrcFromPypi =
         {
@@ -68,182 +46,129 @@ let
         };
     in
     {
-      "3.9-x86_64-linux" = getSrcFromPypi {
-        platform = "manylinux2014_x86_64";
-        dist = "cp39";
-        hash = "sha256-Slbr8FtKTBeRaZ2HTgcvP4CPCYa0AQsU+1SaackMqdw=";
-      };
-      "3.9-aarch64-darwin" = getSrcFromPypi {
-        platform = "macosx_11_0_arm64";
-        dist = "cp39";
-        hash = "sha256-sBVi7IrXVxm30DiXUkiel+trTctMjBE75JFjTVKCrTw=";
-      };
-      "3.9-x86_64-darwin" = getSrcFromPypi {
-        platform = "macosx_10_14_x86_64";
-        dist = "cp39";
-        hash = "sha256-T5jMg3srbG3P4Kt/+esQkxSSCUYRmqOvn6oTlxj/J4c=";
-      };
-
       "3.10-x86_64-linux" = getSrcFromPypi {
         platform = "manylinux2014_x86_64";
         dist = "cp310";
-        hash = "sha256-47zcb45g+FVPQVwU2TATTmAuPKM8OOVGJ0/VRfh1dps=";
+        hash = "sha256-Ya7MuaJ8Z/24RQ9jVyQAGc1FEcudYqROR2R1bThIU60=";
+      };
+      "3.10-aarch64-linux" = getSrcFromPypi {
+        platform = "manylinux2014_aarch64";
+        dist = "cp310";
+        hash = "sha256-7hnBY6j98IOdTBi4il+/tOcxunxDdBbT5Ug+Vwu3ZOQ=";
       };
       "3.10-aarch64-darwin" = getSrcFromPypi {
         platform = "macosx_11_0_arm64";
         dist = "cp310";
-        hash = "sha256-8Djmi9ENGjVUcisLvjbmpEg4RDenWqnSg/aW8O2fjAk=";
+        hash = "sha256-MLL1LLUNdHNK8vR3wlM6elg+O7eyyKzes2Hud9lAV3o=";
       };
       "3.10-x86_64-darwin" = getSrcFromPypi {
         platform = "macosx_10_14_x86_64";
         dist = "cp310";
-        hash = "sha256-pCHSN/jCXShQFm0zRgPGc925tsJvUrxJZwS4eCKXvWY=";
+        hash = "sha256-VcGbnT8zpvxZ9kSqWiH7oCY5zN13bLSptVJmJfV4Of8=";
       };
 
       "3.11-x86_64-linux" = getSrcFromPypi {
         platform = "manylinux2014_x86_64";
         dist = "cp311";
-        hash = "sha256-Rc4PPIQM/4I2z/JsN/Jsn/B4aV+T4MFiwyDCgfUEEnU=";
+        hash = "sha256-J1H/cDfWqZfQvg53zEvjgcWp+buLMU7bdVwTpv2Wn0U=";
+      };
+      "3.11-aarch64-linux" = getSrcFromPypi {
+        platform = "manylinux2014_aarch64";
+        dist = "cp311";
+        hash = "sha256-Q9tYxMQnYnKWNmpWwQMY4fAPUDaQ4X+Uu0NEKT4ZleA=";
       };
       "3.11-aarch64-darwin" = getSrcFromPypi {
         platform = "macosx_11_0_arm64";
         dist = "cp311";
-        hash = "sha256-eThX+vN/Nxyv51L+pfyBH0NeQ7j7S1AgWERKf17M+Ck=";
+        hash = "sha256-P7DqrnNpFXr+y+rVCq8p5z/936d6IzXXIb2XlPPFEOQ=";
       };
       "3.11-x86_64-darwin" = getSrcFromPypi {
         platform = "macosx_10_14_x86_64";
         dist = "cp311";
-        hash = "sha256-L/gpDtx7ksfq5SUX9lSSYz4mey6QZ7rT5MMj0hPnfPU=";
+        hash = "sha256-tn/eq9bf7Qi3do873/tSEWAIX4MFZpvRl77vYdCN4Is=";
       };
 
       "3.12-x86_64-linux" = getSrcFromPypi {
         platform = "manylinux2014_x86_64";
         dist = "cp312";
-        hash = "sha256-RqGqhX9P7uikP8upXA4Kti1AwmzJcwtsaWVZCLo1n40=";
+        hash = "sha256-2tbAqWVnwG0IPARp/sQPIBIQsJk2W9aYvjGm0uyI/Vk=";
+      };
+      "3.12-aarch64-linux" = getSrcFromPypi {
+        platform = "manylinux2014_aarch64";
+        dist = "cp312";
+        hash = "sha256-SW9FsOABojQTCc0MdK8LZwU33O15wWjLIwz8x3PwqoY";
       };
       "3.12-aarch64-darwin" = getSrcFromPypi {
         platform = "macosx_11_0_arm64";
         dist = "cp312";
-        hash = "sha256-jdi//jhTcC9jzZJNoO4lc0pNGc1ckmvgM9dyun0cF10=";
+        hash = "sha256-8zvK/jLJelYuz2iU18QWdMgMCs3t+lQj1Jr1EUcUmHQ=";
       };
       "3.12-x86_64-darwin" = getSrcFromPypi {
         platform = "macosx_10_14_x86_64";
         dist = "cp312";
-        hash = "sha256-1sCaVFMpciRhrwVuc1FG0sjHTCKsdCaoRetp8ya096A=";
+        hash = "sha256-P+/qmF8EFYFvO7r9PwOkNwUCde+brJpywTFOFkSsV8E=";
+      };
+
+      "3.13-x86_64-linux" = getSrcFromPypi {
+        platform = "manylinux2014_x86_64";
+        dist = "cp313";
+        hash = "sha256-LOd7qM2pJZpLypevwcci5CkabEY6Y/jTcsbtyFEX1iU=";
+      };
+      "3.13-aarch64-linux" = getSrcFromPypi {
+        platform = "manylinux2014_aarch64";
+        dist = "cp313";
+        hash = "sha256-JIzKN3Hr8ksHD0lwE2TOraM+YTlEWwbHgsylrFrZK/Q=";
+      };
+      "3.13-aarch64-darwin" = getSrcFromPypi {
+        platform = "macosx_11_0_arm64";
+        dist = "cp313";
+        hash = "sha256-b+MmuK82Y4fdR8zzElg7Kxf+0ScSybdKZIsYoTy9ur8=";
+      };
+      "3.13-x86_64-darwin" = getSrcFromPypi {
+        platform = "macosx_10_14_x86_64";
+        dist = "cp313";
+        hash = "sha256-QeVa5YGKiC5XiehI9vFmh6wTK8+7Wl+hFKXRi3jQXy0=";
       };
     };
-
-  # Note that the prebuilt jaxlib binary requires specific version of CUDA to
-  # work. The cuda12 jaxlib binaries only works with CUDA 12.2, and cuda11
-  # jaxlib binaries only works with CUDA 11.8. This is why we need to find a
-  # binary that matches the provided cudaVersion.
-  gpuSrcVersionString = "cuda${cudaVersion}-${pythonVersion}";
-
-  # 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.
-  gpuSrcs = {
-    "cuda12.2-3.9" = fetchurl {
-      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl";
-      hash = "sha256-d8LIl22gIvmWfoyKfXKElZJXicPQIZxdS4HumhwQGCw=";
-    };
-    "cuda12.2-3.10" = fetchurl {
-      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl";
-      hash = "sha256-PXtWv+UEcMWF8LhWe6Z1UGkf14PG3dkJ0Iop0LiimnQ=";
-    };
-    "cuda12.2-3.11" = fetchurl {
-      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl";
-      hash = "sha256-QO2WSOzmJ48VaCha596mELiOfPsAGLpGctmdzcCHE/o=";
-    };
-    "cuda12.2-3.12" = fetchurl {
-      url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl";
-      hash = "sha256-ixWMaIChy4Ammsn23/3cCoala0lFibuUxyUr3tjfFKU=";
-    };
-  };
 in
 buildPythonPackage {
   pname = "jaxlib";
   inherit version;
   format = "wheel";
 
-  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."${pythonVersion}-${stdenv.hostPlatform.system}"
-        or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
-      )
-    else
-      gpuSrcs."${gpuSrcVersionString}" or (throw "jaxlib-bin: No gpuSrc for ${gpuSrcVersionString}");
+  src = (
+    srcs."${pythonVersion}-${stdenv.hostPlatform.system}"
+      or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
+  );
 
   # Prebuilt wheels are dynamically linked against things that nix can't find.
   # Run `autoPatchelfHook` to automagically fix them.
-  nativeBuildInputs =
-    lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ]
-    ++ lib.optionals cudaSupport [ autoAddDriverRunpath ];
+  nativeBuildInputs = lib.optionals stdenv.hostPlatform.isLinux [ autoPatchelfHook ];
   # Dynamic link dependencies
   buildInputs = [ (lib.getLib stdenv.cc.cc) ];
 
-  # 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 = lib.optional cudaSupport ''
-    shopt -s globstar
-
-    for file in $out/**/*.so; do
-      patchelf --add-rpath "${cudaLibPath}" "$file"
-    done
-  '';
-
-  propagatedBuildInputs = [
+  dependencies = [
     absl-py
     flatbuffers
     ml-dtypes
     scipy
   ];
 
-  # 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 = lib.optional cudaSupport ''
-    mkdir -p $out/${python.sitePackages}/jaxlib/cuda/bin
-    ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jaxlib/cuda/bin/ptxas
-  '';
+  pythonImportsCheck = [ "jaxlib" ];
 
-  inherit (jaxlib-build) pythonImportsCheck;
-
-  meta = with lib; {
+  meta = {
     description = "Prebuilt jaxlib backend from PyPi";
     homepage = "https://github.com/google/jax";
-    sourceProvenance = with sourceTypes; [ binaryNativeCode ];
-    license = licenses.asl20;
-    maintainers = with maintainers; [ samuela ];
-    platforms = [
-      "aarch64-darwin"
-      "x86_64-linux"
-      "x86_64-darwin"
-    ];
-    broken =
-      !(cudaSupport -> lib.versionAtLeast cudaVersion "11.1")
-      || !(cudaSupport -> lib.versionAtLeast cudaPackages.cudnn.version "8.2")
-      || !(cudaSupport -> stdenv.hostPlatform.isLinux)
-      || !(cudaSupport -> (gpuSrcs ? "cuda${cudaVersion}-${pythonVersion}"))
+    sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ samuela ];
+    badPlatforms = [
       # Fails at pythonImportsCheckPhase:
       # ...-python-imports-check-hook.sh/nix-support/setup-hook: line 10: 28017 Illegal instruction: 4
       # /nix/store/5qpssbvkzfh73xih07xgmpkj5r565975-python3-3.11.9/bin/python3.11 -c
       # 'import os; import importlib; list(map(lambda mod: importlib.import_module(mod), os.environ["pythonImportsCheck"].split()))'
-      || (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64);
+      "x86_64-darwin"
+    ];
   };
 }
diff --git a/pkgs/development/python-modules/jaxlib/prefetch.sh b/pkgs/development/python-modules/jaxlib/prefetch.sh
index 46b2573e3c30..22c5253ab0d8 100755
--- a/pkgs/development/python-modules/jaxlib/prefetch.sh
+++ b/pkgs/development/python-modules/jaxlib/prefetch.sh
@@ -2,17 +2,23 @@
 #! nix-shell -i sh -p jq
 
 prefetch () {
-    expr="(import <nixpkgs> { system = \"$2\"; config.cudaSupport = $3; }).python$1.pkgs.jaxlib-bin.src.url"
+    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 hash to-sri --type sha256 "$sha256"
+    nix hash convert --to sri --hash-algo sha256 "$sha256"
     echo
 }
 
-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"
+for py in "310" "311" "312" "313"; do
+    prefetch "$py" "x86_64-linux" "jaxlib-bin"
+    prefetch "$py" "aarch64-linux" "jaxlib-bin"
+    prefetch "$py" "aarch64-darwin" "jaxlib-bin"
+    prefetch "$py" "x86_64-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
diff --git a/pkgs/development/python-modules/nifty8/default.nix b/pkgs/development/python-modules/nifty8/default.nix
index ba3e0f14932e..70e9140123db 100644
--- a/pkgs/development/python-modules/nifty8/default.nix
+++ b/pkgs/development/python-modules/nifty8/default.nix
@@ -1,12 +1,9 @@
 {
-  config,
   lib,
   stdenv,
   buildPythonPackage,
   fetchFromGitLab,
 
-  cudaSupport ? config.cudaSupport,
-
   # build-system
   setuptools,
 
@@ -48,7 +45,7 @@ buildPythonPackage rec {
     ducc0
     h5py
     jax
-    (jaxlib.override { inherit cudaSupport; })
+    jaxlib
     matplotlib
     mpi4py
     mpi
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 1056a924484e..d44321e85813 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -6560,11 +6560,13 @@ self: super: with self; {
 
   jax = callPackage ../development/python-modules/jax { };
 
+  jax-cuda12-pjrt = callPackage ../development/python-modules/jax-cuda12-pjrt { };
+
+  jax-cuda12-plugin = callPackage ../development/python-modules/jax-cuda12-plugin { };
+
   jax-jumpy = callPackage ../development/python-modules/jax-jumpy { };
 
-  jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix {
-    inherit (pkgs.config) cudaSupport;
-  };
+  jaxlib-bin = callPackage ../development/python-modules/jaxlib/bin.nix { };
 
   jaxlib-build = callPackage ../development/python-modules/jaxlib rec {
     # Some platforms don't have `cudaSupport` defined, hence the need for 'or false'.
@@ -6572,16 +6574,7 @@ self: super: with self; {
     IOKit = pkgs.darwin.apple_sdk_11_0.IOKit;
   };
 
-  # Use the -bin on macOS since the source build doesn't support it (see #323154)
-  jaxlib = if jaxlib-build.meta.unsupported then jaxlib-bin else jaxlib-build;
-
-  jaxlibWithCuda = self.jaxlib.override {
-    cudaSupport = true;
-  };
-
-  jaxlibWithoutCuda = self.jaxlib.override {
-    cudaSupport = false;
-  };
+  jaxlib = jaxlib-bin;
 
   jaxopt = callPackage ../development/python-modules/jaxopt { };