Something went wrong. Try again.
What happens if I push all of nixpkgs
Something went wrong. Try again.
2.9 kB · 129 lines
Nix
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130{ lib, stdenv, absl-py, buildPythonPackage, fetchFromGitHub,
# build-system flit-core,
# dependencies aiofiles, etils, humanize, importlib-resources, jax, msgpack, nest-asyncio, numpy, protobuf, pyyaml, simplejson, tensorstore, typing-extensions,
# tests chex, google-cloud-logging, mock, optax, portpicker, pytest-xdist, pytestCheckHook, safetensors,}:
buildPythonPackage rec { pname = "orbax-checkpoint"; version = "0.11.20"; pyproject = true;
src = fetchFromGitHub { owner = "google"; repo = "orbax"; tag = "v${version}"; hash = "sha256-ZZTmWW0PZHk8Evhk6JnbAyl7RQOAxN4GwXUv5M6GfIQ="; };
sourceRoot = "${src.name}/checkpoint";
build-system = [ flit-core ];
pythonRelaxDeps = [ "jax" ];
dependencies = [ absl-py aiofiles etils humanize importlib-resources jax msgpack nest-asyncio numpy protobuf pyyaml simplejson tensorstore typing-extensions ];
nativeCheckInputs = [ chex google-cloud-logging mock optax portpicker pytest-xdist pytestCheckHook safetensors ];
pythonImportsCheck = [ "orbax" "orbax.checkpoint" ];
disabledTests = [ # Flaky # AssertionError: 2 not greater than 2.0046136379241943 "test_async_mkdir_parallel" "test_async_mkdir_sequential" ] ++ lib.optionals stdenv.hostPlatform.isDarwin [ # Probably failing because of a filesystem impurity # self.assertFalse(os.path.exists(dst_dir)) # AssertionError: True is not false "test_create_snapshot" ];
disabledTestPaths = [ # E absl.flags._exceptions.DuplicateFlagError: The flag 'num_processes' is defined twice. # First from multiprocess_test, Second from orbax.checkpoint._src.testing.multiprocess_test. # Description from first occurrence: Number of processes to use. # https://github.com/google/orbax/issues/1580 "orbax/checkpoint/experimental/emergency/"
# Circular dependency flax "orbax/checkpoint/_src/metadata/empty_values_test.py" "orbax/checkpoint/_src/metadata/tree_rich_types_test.py" "orbax/checkpoint/_src/metadata/tree_test.py" "orbax/checkpoint/_src/testing/test_tree_utils.py" "orbax/checkpoint/_src/tree/parts_of_test.py" "orbax/checkpoint/_src/tree/structure_utils_test.py" "orbax/checkpoint/_src/tree/utils_test.py" "orbax/checkpoint/single_host_test.py" "orbax/checkpoint/transform_utils_test.py" ];
meta = { 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/v${version}/checkpoint/CHANGELOG.md"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ fab ]; };}