diff --git a/pkgs/development/python-modules/jaxlib/default.nix b/pkgs/development/python-modules/jaxlib/default.nix
index 1f4c8a9693b..2bcf5e9c17c 100644
--- a/pkgs/development/python-modules/jaxlib/default.nix
+++ b/pkgs/development/python-modules/jaxlib/default.nix
@@ -63,7 +63,7 @@ let
     # aarch64-darwin is broken because of https://github.com/bazelbuild/rules_cc/pull/136
     # however even with that fix applied, it doesn't work for everyone:
     # https://github.com/NixOS/nixpkgs/pull/184395#issuecomment-1207287129
-    broken = stdenv.isAarch64 || stdenv.isDarwin;
+    broken = stdenv.isDarwin;
   };
 
   cudatoolkit_joined = symlinkJoin {
@@ -129,6 +129,11 @@ let
     "zlib"
   ];
 
+  arch =
+    # KeyError: ('Linux', 'arm64')
+    if stdenv.targetPlatform.isLinux && stdenv.targetPlatform.linuxArch == "arm64" then "aarch64"
+    else stdenv.targetPlatform.linuxArch;
+
   bazel-build = buildBazelPackage rec {
     name = "bazel-build-${pname}-${version}";
 
@@ -291,7 +296,7 @@ let
       '' else throw "Unsupported stdenv.cc: ${stdenv.cc}");
 
       installPhase = ''
-        ./bazel-bin/build/build_wheel --output_path=$out --cpu=${stdenv.targetPlatform.linuxArch}
+        ./bazel-bin/build/build_wheel --output_path=$out --cpu=${arch}
       '';
     };
 
@@ -299,11 +304,11 @@ let
   };
   platformTag =
     if stdenv.targetPlatform.isLinux then
-      "manylinux2014_${stdenv.targetPlatform.linuxArch}"
+      "manylinux2014_${arch}"
     else if stdenv.system == "x86_64-darwin" then
-      "macosx_10_9_${stdenv.targetPlatform.linuxArch}"
+      "macosx_10_9_${arch}"
     else if stdenv.system == "aarch64-darwin" then
-      "macosx_11_0_${stdenv.targetPlatform.linuxArch}"
+      "macosx_11_0_${arch}"
     else throw "Unsupported target platform: ${stdenv.targetPlatform}";
 
 in