From df9a844ef7680f0f753c4170efe147fb3f1a5ac6 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 18 Apr 2023 17:39:08 +0200
Subject: [PATCH 1/2] python3Packages.jaxtyping: init at 0.2.20

---
 .../python-modules/jaxtyping/default.nix      | 64 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 66 insertions(+)
 create mode 100644 pkgs/development/python-modules/jaxtyping/default.nix

diff --git a/pkgs/development/python-modules/jaxtyping/default.nix b/pkgs/development/python-modules/jaxtyping/default.nix
new file mode 100644
index 000000000000..42071885eb60
--- /dev/null
+++ b/pkgs/development/python-modules/jaxtyping/default.nix
@@ -0,0 +1,64 @@
+{ lib
+, buildPythonPackage
+, fetchFromGitHub
+, hatchling
+, numpy
+, typeguard
+, typing-extensions
+, cloudpickle
+, equinox
+, jax
+, jaxlib
+, torch
+, pytestCheckHook
+}:
+
+let
+  jaxtyping = buildPythonPackage rec {
+    pname = "jaxtyping";
+    version = "0.2.20";
+    format = "pyproject";
+
+    src = fetchFromGitHub {
+      owner = "google";
+      repo = pname;
+      rev = "refs/tags/v${version}";
+      hash = "sha256-q/KQGV7I7w5p7VP8C9BDUHfPsuCMf2v304qiH+XCzyU=";
+    };
+
+    nativeBuildInputs = [
+      hatchling
+    ];
+
+    propagatedBuildInputs = [
+      numpy
+      typeguard
+      typing-extensions
+    ];
+
+    nativeCheckInputs = [
+      cloudpickle
+      equinox
+      jax
+      jaxlib
+      pytestCheckHook
+      torch
+    ];
+
+    doCheck = false;
+
+    # Enable tests via passthru to avoid cyclic dependency with equinox.
+    passthru.tests = {
+      check = jaxtyping.overridePythonAttrs { doCheck = true; };
+    };
+
+    pythonImportsCheck = [ "jaxtyping" ];
+
+    meta = with lib; {
+      description = "Type annotations and runtime checking for JAX arrays and PyTrees";
+      homepage = "https://github.com/google/jaxtyping";
+      license = licenses.mit;
+      maintainers = with maintainers; [ GaetanLepage ];
+    };
+  };
+ in jaxtyping
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index d4bb7f8cfc43..8ad04627a0a2 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -5336,6 +5336,8 @@ self: super: with self; {
 
   jaxopt = callPackage ../development/python-modules/jaxopt { };
 
+  jaxtyping = callPackage ../development/python-modules/jaxtyping { };
+
   jaydebeapi = callPackage ../development/python-modules/jaydebeapi { };
 
   jc = callPackage ../development/python-modules/jc { };

From 644078614c69689daad44134b131a81d8b7195e3 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Tue, 18 Apr 2023 17:39:43 +0200
Subject: [PATCH 2/2] python3Packages.equinox: init at 0.10.11

---
 .../python-modules/equinox/default.nix        | 49 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 51 insertions(+)
 create mode 100644 pkgs/development/python-modules/equinox/default.nix

diff --git a/pkgs/development/python-modules/equinox/default.nix b/pkgs/development/python-modules/equinox/default.nix
new file mode 100644
index 000000000000..853e239500e2
--- /dev/null
+++ b/pkgs/development/python-modules/equinox/default.nix
@@ -0,0 +1,49 @@
+{ lib
+, buildPythonPackage
+, fetchFromGitHub
+, hatchling
+, jax
+, jaxlib
+, jaxtyping
+, typing-extensions
+, beartype
+, pytestCheckHook
+}:
+
+buildPythonPackage rec {
+  pname = "equinox";
+  version = "0.10.11";
+  format = "pyproject";
+
+  src = fetchFromGitHub {
+    owner = "patrick-kidger";
+    repo = pname;
+    rev = "refs/tags/v${version}";
+    hash = "sha256-JffuPplIROPog29FBsWH9cQHSkrFKuXjaTjjEwIqW/0=";
+  };
+
+  nativeBuildInputs = [
+    hatchling
+  ];
+
+  propagatedBuildInputs = [
+    jax
+    jaxlib
+    jaxtyping
+    typing-extensions
+  ];
+
+  nativeCheckInputs = [
+    beartype
+    pytestCheckHook
+  ];
+
+  pythonImportsCheck = [ "equinox" ];
+
+  meta = with lib; {
+    description = "A JAX library based around a simple idea: represent parameterised functions (such as neural networks) as PyTrees";
+    homepage = "https://github.com/patrick-kidger/equinox";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ GaetanLepage ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 8ad04627a0a2..125728796420 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -3406,6 +3406,8 @@ self: super: with self; {
 
   epson-projector = callPackage ../development/python-modules/epson-projector { };
 
+  equinox = callPackage ../development/python-modules/equinox { };
+
   eradicate = callPackage ../development/python-modules/eradicate { };
 
   es-client = callPackage ../development/python-modules/es-client { };