diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
index 628bd947d071..7e5dd0648d48 100644
--- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
+++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix
@@ -14,6 +14,18 @@
 let
   inherit (jaxlib) version;
 
+  platforms = {
+    x86_64-linux = {
+      name = "manylinux_2_27_x86_64";
+      hash = "sha256-/dnXJAjhf3ojvfombpyQE2YgpyTfWEYp81BnIuKy9DQ=";
+    };
+    aarch64-linux = {
+      name = "manylinux_2_27_aarch64";
+      hash = "sha256-QAIdRVvZRSJwTjIdMOitomp8Cm75XEp1QqvaATiinRw=";
+    };
+  };
+  currentPlatform = platforms.${stdenv.hostPlatform.system};
+
   cudaLibPath = lib.makeLibraryPath (
     with cudaPackages;
     [
@@ -43,18 +55,8 @@ buildPythonPackage (finalAttrs: {
     format = "wheel";
     python = "py3";
     dist = "py3";
-    platform =
-      {
-        x86_64-linux = "manylinux_2_27_x86_64";
-        aarch64-linux = "manylinux_2_27_aarch64";
-      }
-      .${stdenv.hostPlatform.system};
-    hash =
-      {
-        x86_64-linux = "sha256-/dnXJAjhf3ojvfombpyQE2YgpyTfWEYp81BnIuKy9DQ=";
-        aarch64-linux = "sha256-QAIdRVvZRSJwTjIdMOitomp8Cm75XEp1QqvaATiinRw=";
-      }
-      .${stdenv.hostPlatform.system};
+    platform = currentPlatform.name;
+    inherit (currentPlatform) hash;
   };
 
   nativeBuildInputs = [
@@ -100,11 +102,7 @@ buildPythonPackage (finalAttrs: {
     sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
     license = lib.licenses.asl20;
     maintainers = with lib.maintainers; [ natsukium ];
-    # Keep in sync with `src` above
-    platforms = [
-      "x86_64-linux"
-      "aarch64-linux"
-    ];
+    platforms = lib.attrNames platforms;
     # see CUDA compatibility matrix
     # https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder
     broken = !(lib.versionAtLeast cudaPackages.cudnn.version "9.1");
diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
index c7bf3db79514..02a949297aec 100644
--- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
+++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix
@@ -15,57 +15,29 @@ let
   inherit (jaxlib) version;
   inherit (jax-cuda12-pjrt) cudaLibPath;
 
-  getSrcFromPypi =
-    {
-      platform,
-      dist,
-      hash,
-    }:
-    fetchPypi {
-      inherit
-        version
-        platform
-        dist
-        hash
-        ;
-      pname = "jax_cuda12_plugin";
-      format = "wheel";
-      python = dist;
-      abi = dist;
+  platforms = {
+    x86_64-linux = {
+      name = "manylinux_2_27_x86_64";
+      hashes = {
+        cp312 = "sha256-8BVPvyNrsiMab8pSBN0eQ6wMHPDeOgAg//kQPlP7hXI=";
+        cp313 = "sha256-2O/FH9SZiimQxGfmPDhs/ckYfH0QWCXFxKt0Pn1SLRE=";
+        cp314 = "sha256-a7jxBoIVF+NEiKUyDkX13vSkYo8qJUr6bix+7DVBm74=";
+      };
     };
-
-  srcs = {
-    "3.12-x86_64-linux" = getSrcFromPypi {
-      platform = "manylinux_2_27_x86_64";
-      dist = "cp312";
-      hash = "sha256-8BVPvyNrsiMab8pSBN0eQ6wMHPDeOgAg//kQPlP7hXI=";
-    };
-    "3.12-aarch64-linux" = getSrcFromPypi {
-      platform = "manylinux_2_27_aarch64";
-      dist = "cp312";
-      hash = "sha256-8qOC5dv2YzB0C/TbOQPGEuhLJN4JkTx5U+0kA8KWtoY=";
-    };
-    "3.13-x86_64-linux" = getSrcFromPypi {
-      platform = "manylinux_2_27_x86_64";
-      dist = "cp313";
-      hash = "sha256-2O/FH9SZiimQxGfmPDhs/ckYfH0QWCXFxKt0Pn1SLRE=";
-    };
-    "3.13-aarch64-linux" = getSrcFromPypi {
-      platform = "manylinux_2_27_aarch64";
-      dist = "cp313";
-      hash = "sha256-6KjWb0PQCzXAQwNmsRfBthKDdk8NENOWWr/jn4coSG8=";
-    };
-    "3.14-x86_64-linux" = getSrcFromPypi {
-      platform = "manylinux_2_27_x86_64";
-      dist = "cp314";
-      hash = "sha256-a7jxBoIVF+NEiKUyDkX13vSkYo8qJUr6bix+7DVBm74=";
-    };
-    "3.14-aarch64-linux" = getSrcFromPypi {
-      platform = "manylinux_2_27_aarch64";
-      dist = "cp314";
-      hash = "sha256-PKEZswHZqTvLg3XGk3z3lJL/WVHKXLuLRvfjSvQtz9M=";
+    aarch64-linux = {
+      name = "manylinux_2_27_aarch64";
+      hashes = {
+        cp312 = "sha256-8qOC5dv2YzB0C/TbOQPGEuhLJN4JkTx5U+0kA8KWtoY=";
+        cp313 = "sha256-6KjWb0PQCzXAQwNmsRfBthKDdk8NENOWWr/jn4coSG8=";
+        cp314 = "sha256-PKEZswHZqTvLg3XGk3z3lJL/WVHKXLuLRvfjSvQtz9M=";
+      };
     };
   };
+  currentPlatform =
+    platforms.${stdenv.hostPlatform.system}
+      or (throw "jax-cuda12-plugin is not supported on ${stdenv.hostPlatform.system}");
+
+  dist = "cp${lib.replaceStrings [ "." ] [ "" ] python.pythonVersion}";
 in
 buildPythonPackage {
   pname = "jax-cuda12-plugin";
@@ -73,10 +45,17 @@ buildPythonPackage {
   pyproject = false;
   __structuredAttrs = true;
 
-  src = (
-    srcs."${python.pythonVersion}-${stdenv.hostPlatform.system}"
-      or (throw "python${python.pythonVersion}Packages.jax-cuda12-plugin is not supported on ${stdenv.hostPlatform.system}")
-  );
+  src = fetchPypi {
+    pname = "jax_cuda12_plugin";
+    inherit version dist;
+    format = "wheel";
+    python = dist;
+    abi = dist;
+    platform = currentPlatform.name;
+    hash =
+      currentPlatform.hashes.${dist}
+        or (throw "python${python.pythonVersion}Packages.jax-cuda12-plugin is not supported");
+  };
 
   nativeBuildInputs = [
     autoPatchelfHook
@@ -119,11 +98,7 @@ buildPythonPackage {
     sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
     license = lib.licenses.asl20;
     maintainers = with lib.maintainers; [ natsukium ];
-    # Keep in sync with `srcs` above
-    platforms = [
-      "x86_64-linux"
-      "aarch64-linux"
-    ];
+    platforms = lib.attrNames platforms;
     # see CUDA compatibility matrix
     # https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder
     broken = !(lib.versionAtLeast cudaPackages.cudnn.version "9.1");