Merge pull request #247925 from GaetanLepage/equinox

python3Packages.equinox: init at 0.10.11

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 ];
+  };
+}

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

pkgs/top-level/python-packages.nix

@@ -3410,6 +3410,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 { };
@@ -5340,6 +5342,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 { };