python3Packages.orbax-export: init at 0.0.8

pkgs/development/python-modules/orbax-export/default.nix

@@ -0,0 +1,90 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchPypi,
+
+  # build-system
+  flit-core,
+
+  # dependencies
+  absl-py,
+  dataclasses-json,
+  etils,
+  jax,
+  jaxlib,
+  jaxtyping,
+  numpy,
+  orbax-checkpoint,
+  protobuf,
+  tensorflow,
+
+  # tests
+  chex,
+  flax,
+  pytestCheckHook,
+
+  # passthru
+  orbax-export,
+}:
+
+buildPythonPackage (finalAttrs: {
+  pname = "orbax-export";
+  version = "0.0.8";
+  pyproject = true;
+
+  # Tags on the GitHub repo don't match the Pypi releases for orbax-export
+  src = fetchPypi {
+    pname = "orbax_export";
+    inherit (finalAttrs) version;
+    hash = "sha256-VE7vVk4qbxfNEbEWf+vjSLe3z1bZV13plKM9VhPdVoo=";
+  };
+
+  build-system = [
+    flit-core
+  ];
+
+  dependencies = [
+    absl-py
+    dataclasses-json
+    etils
+    jax
+    jaxlib
+    jaxtyping
+    numpy
+    orbax-checkpoint
+    protobuf
+    tensorflow
+  ];
+
+  pythonImportsCheck = [
+    "orbax"
+    "orbax.export"
+    "orbax.export.bfloat16_toolkit.python"
+  ];
+
+  nativeCheckInputs = [
+    chex
+    flax
+    pytestCheckHook
+  ];
+
+  preCheck = ''
+    cd orbax/export
+    rm -rf ./**/__init__.py
+    rm -rf typing
+  '';
+
+  # Circular dependency with flax
+  doCheck = false;
+
+  passthru.tests.pytest = orbax-export.overridePythonAttrs {
+    doCheck = true;
+  };
+
+  meta = {
+    description = "Serialization library for JAX users, enabling the exporting of JAX models to the TensorFlow SavedModel format";
+    homepage = "https://github.com/google/orbax/tree/main/export";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+})

pkgs/top-level/python-packages.nix

@@ -11823,6 +11823,8 @@ self: super: with self; {
 
   orbax-checkpoint = callPackage ../development/python-modules/orbax-checkpoint { };
 
+  orbax-export = callPackage ../development/python-modules/orbax-export { };
+
   ordered-set = callPackage ../development/python-modules/ordered-set { };
 
   orderedmultidict = callPackage ../development/python-modules/orderedmultidict { };