1 files changed, 78 insertions, 0 deletions
diff --git a/nixpkgs/pkgs/development/python-modules/orbax-checkpoint/default.nix b/nixpkgs/pkgs/development/python-modules/orbax-checkpoint/default.nix
new file mode 100644
index 000000000000..0f9d467335ce
--- /dev/null
+++ b/nixpkgs/pkgs/development/python-modules/orbax-checkpoint/default.nix
@@ -0,0 +1,78 @@
+{ lib
+, absl-py
+, buildPythonPackage
+, cached-property
+, etils
+, fetchPypi
+, flit-core
+, importlib-resources
+, jax
+, jaxlib
+, msgpack
+, nest-asyncio
+, numpy
+, protobuf
+, pytest-xdist
+, pytestCheckHook
+, pythonOlder
+, pyyaml
+, tensorstore
+, typing-extensions
+}:
+
+buildPythonPackage rec {
+  pname = "orbax-checkpoint";
+  version = "0.5.3";
+  pyproject = true;
+
+  disabled = pythonOlder "3.9";
+
+  src = fetchPypi {
+    pname = "orbax_checkpoint";
+    inherit version;
+    hash = "sha256-FXKQTLv+hROSfg2A+AtzDg7y9oAzLTwoENhENTKTi0U=";
+  };
+
+  nativeBuildInputs = [
+    flit-core
+  ];
+
+  propagatedBuildInputs = [
+    absl-py
+    cached-property
+    etils
+    importlib-resources
+    jax
+    jaxlib
+    msgpack
+    nest-asyncio
+    numpy
+    protobuf
+    pyyaml
+    tensorstore
+    typing-extensions
+  ];
+
+  nativeCheckInputs = [
+    pytest-xdist
+    pytestCheckHook
+  ];
+
+  pythonImportsCheck = [
+    "orbax"
+  ];
+
+  disabledTestPaths = [
+    # Circular dependency flax
+    "orbax/checkpoint/transform_utils_test.py"
+    "orbax/checkpoint/utils_test.py"
+  ];
+
+  meta = with lib; {
+    description = "Orbax provides common utility libraries for JAX users";
+    homepage = "https://github.com/google/orbax/tree/main/checkpoint";
+    changelog = "https://github.com/google/orbax/blob/${version}/CHANGELOG.md";
+    license = licenses.asl20;
+    maintainers = with maintainers; [fab ];
+  };
+}