diff --git a/pkgs/development/python-modules/jaxtyping/default.nix b/pkgs/development/python-modules/jaxtyping/default.nix
index 1193fec64749..8390583865b8 100644
--- a/pkgs/development/python-modules/jaxtyping/default.nix
+++ b/pkgs/development/python-modules/jaxtyping/default.nix
@@ -18,58 +18,58 @@
   pytestCheckHook,
   tensorflow,
   torch,
+
+  # passthru
+  jaxtyping,
 }:
 
-let
-  self = buildPythonPackage rec {
-    pname = "jaxtyping";
-    version = "0.3.4";
-    pyproject = true;
+buildPythonPackage (finalAttrs: {
+  pname = "jaxtyping";
+  version = "0.3.5";
+  pyproject = true;
 
-    src = fetchFromGitHub {
-      owner = "google";
-      repo = "jaxtyping";
-      tag = "v${version}";
-      hash = "sha256-u9JsZm5uz2XXQnhgRe5P/fJrKKnyIMdCTnBLCbW+kAw=";
-    };
+  src = fetchFromGitHub {
+    owner = "patrick-kidger";
+    repo = "jaxtyping";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-0Vt6UD1xQkwve6yDVi5XQCoJ/IsJWHCkGesj66myQq4=";
+  };
 
-    build-system = [ hatchling ];
+  build-system = [ hatchling ];
 
-    dependencies = [
-      wadler-lindig
-    ];
+  dependencies = [
+    wadler-lindig
+  ];
 
-    pythonImportsCheck = [ "jaxtyping" ];
+  pythonImportsCheck = [ "jaxtyping" ];
 
-    nativeCheckInputs = [
-      cloudpickle
-      equinox
-      ipython
-      jax
-      jaxlib
-      pytestCheckHook
-      tensorflow
-      torch
-    ];
+  nativeCheckInputs = [
+    cloudpickle
+    equinox
+    ipython
+    jax
+    jaxlib
+    pytestCheckHook
+    tensorflow
+    torch
+  ];
 
-    doCheck = false;
+  doCheck = false;
 
-    # Enable tests via passthru to avoid cyclic dependency with equinox.
-    passthru.tests = {
-      check = self.overridePythonAttrs {
-        # We disable tests because they complain about the version of typeguard being too new.
-        doCheck = false;
-        catchConflicts = false;
-      };
-    };
-
-    meta = {
-      description = "Type annotations and runtime checking for JAX arrays and PyTrees";
-      homepage = "https://github.com/google/jaxtyping";
-      changelog = "https://github.com/patrick-kidger/jaxtyping/releases/tag/v${version}";
-      license = lib.licenses.mit;
-      maintainers = with lib.maintainers; [ GaetanLepage ];
+  # Enable tests via passthru to avoid cyclic dependency with equinox.
+  passthru.tests = {
+    check = jaxtyping.overridePythonAttrs {
+      # We disable tests because they complain about the version of typeguard being too new.
+      doCheck = false;
+      catchConflicts = false;
     };
   };
-in
-self
+
+  meta = {
+    description = "Type annotations and runtime checking for JAX arrays and PyTrees";
+    homepage = "https://github.com/patrick-kidger/jaxtyping";
+    changelog = "https://github.com/patrick-kidger/jaxtyping/releases/tag/${finalAttrs.src.tag}";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+})