Commits

Branch selector

 main

User selector

All users

Datepicker

All time

Commit History

Commits on Oct 11, 2024

• Modify `enable_pinned_host_transfer` option to be provided once for the entire pytree, since it's not really reasonable to customize this on a per-array level.

  Show description for 8b4e90d
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  8b4e90d

  Copy full SHA for 8b4e90d

• Rename `_single_item` to `_default_item` in `CheckpointManager`.

  Show description for fef8e18
  

  BlaziusMaximus

  authored and

  Orbax Authors

  committed
  fef8e18

  Copy full SHA for fef8e18

Commits on Oct 10, 2024

• Add a `SaveArgs` option that allows disabling pinned host transfer on a per-array basis.

  Show description for 7ffdc88
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  7ffdc88

  Copy full SHA for 7ffdc88

• #cleanup Upgrade type annotations in `tensorstore_utils` to 3.10.

  Show description for 40ae6ca
  

  dicentra13

  authored and

  Orbax Authors

  committed
  40ae6ca

  Copy full SHA for 40ae6ca

• Refactor `build_array_tspec_for_write` into a dataclass.

  Show description for 83a2c7c
  

  dicentra13

  authored and

  Orbax Authors

  committed
  83a2c7c

  Copy full SHA for 83a2c7c

Commits on Oct 8, 2024

• Add monitoring to record how long directory creation takes (per sequential instance).

  Show description for 4bb4362
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  4bb4362

  Copy full SHA for 4bb4362

• Refactor calls to `choose_chunk_shape` to simplify `build_zarr_shard_and_chunk_metadata`.

  Show description for 4c305e6
  

  dicentra13

  authored and

  Orbax Authors

  committed
  4c305e6

  Copy full SHA for 4c305e6

Commits on Oct 7, 2024

• Now that JAX supports donate on device_put, use that instead of the ugly device_put -> block -> delete routine

  Show description for db46fc7
  

  yashk2810

  authored and

  Orbax Authors

  committed
  db46fc7

  Copy full SHA for db46fc7

• Remove orbax-checkpoint support for Python 3.9.

  Show description for d1b8aaa
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  d1b8aaa

  Copy full SHA for d1b8aaa

Commits on Oct 3, 2024

• Refine several API docs and add a test to prevent colab docs from breaking.

  Show description for 410ca56
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  410ca56

  Copy full SHA for 410ca56

• Fix `_determine_single_item_mode_from_directory` by incorporating step name format.

  Show description for 8bf99e3
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  8bf99e3

  Copy full SHA for 8bf99e3

Commits on Oct 2, 2024

• Rename CheckpointMetadata to StepMetadata.

  Show description for 367ad4d
  

  BlaziusMaximus

  authored and

  Orbax Authors

  committed
  367ad4d

  Copy full SHA for 367ad4d

• Allow registering an item in `DefaultCheckpointHandlerRegistry` without providing the actual handler, as long as the provided args correspond to a globally registered handler. This allows for sligh…

  Show description for 07b17c1
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  07b17c1

  Copy full SHA for 07b17c1

• Support `empty tuple` values in checkpointable params.

  Show description for aad5e72
  

  niketkumar

  authored and

  Orbax Authors

  committed
  aad5e72

  Copy full SHA for aad5e72

• Use `CheckpointHandlerRegistry` when `CheckpointManager` is default-constructed, e.g. `CheckpointManager(directory)`. Previously this usage pattern would default to single-item mode, but the new be…

  Show description for a01eab8
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  a01eab8

  Copy full SHA for a01eab8

• No public description

  Show description for e40ba06
  

  Orbax Authors

  committed
  e40ba06

  Copy full SHA for e40ba06

Commits on Sep 30, 2024

• No public description

  Show description for 4279d83
  

  niketkumar

  authored and

  Orbax Authors

  committed
  4279d83

  Copy full SHA for 4279d83

• Make subchunking using `shard_axes` deterministic.

  Show description for 198266d
  

  dicentra13

  authored and

  Orbax Authors

  committed
  198266d

  Copy full SHA for 198266d

Commits on Sep 27, 2024

• Use `__future__.annotations` module to modernize type definitions.

  Show description for 442ffb8
  

  niketkumar

  authored and

  Orbax Authors

  committed
  442ffb8

  Copy full SHA for 442ffb8

• Adjust `CompositeCheckpointHandler` behavior for unregistered items and empty `args`. Now, `metadata` only returns entries for which an item actually exists in the checkpoint. `restore` will raise …

  Show description for 49acec9
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  49acec9

  Copy full SHA for 49acec9

Commits on Sep 26, 2024

• Refactor to extract a separate module, `asyncio_utils`, for asyncio helper functions from `path/async_utils` module.

  Show description for d9c9c3a
  

  niketkumar

  authored and

  Orbax Authors

  committed
  d9c9c3a

  Copy full SHA for d9c9c3a

• Remove logging of ts_spec in validation step - this is very verbose.

  Show description for be1ecdc
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  be1ecdc

  Copy full SHA for be1ecdc

• Fix a few typos in the comments.

  Show description for 10d3d0a
  

  Orbax Authors

  committed
  10d3d0a

  Copy full SHA for 10d3d0a

Commits on Sep 25, 2024

• Adds NotImplemented errors to obm_module and adds related tests to jax_module_test.

  Show description for c62f130
  

  Orbax Authors

  committed
  c62f130

  Copy full SHA for c62f130

• Log error if number of checkpoints exceeds max_to_keep instead of raising error in emergency checkpointing.

  Show description for cd4b736
  

  cpgaffney1

  authored and

  Orbax Authors

  committed
  cd4b736

  Copy full SHA for cd4b736

• Cleans up the duplicated logic in jax_module after moving much of it to tensorflow_module in a previous checkin. Updates the tests for jax_module to focus on its smaller scope. Cleans up some of th…

  Show description for ff9f38f
  

  Orbax Authors

  committed
  ff9f38f

  Copy full SHA for ff9f38f

Commits on Sep 24, 2024

• Moves TF dependent code from jax_module in to tensorflow_module and from jax_module_test to tensorflow_module_test. This code is duplicated for now and will be cleaned up in a subsequent cl.

  Show description for fb7272d
  

  Orbax Authors

  committed
  fb7272d

  Copy full SHA for fb7272d

• Update links and information in Orbax checkpointing documentation.

  Show description for 9da7198
  

  BlaziusMaximus

  authored and

  Orbax Authors

  committed
  9da7198

  Copy full SHA for 9da7198

Commits on Sep 23, 2024

• Fix inaccurate comment

  Show description for aa2bc58
  

  Orbax Authors

  committed
  aa2bc58

  Copy full SHA for aa2bc58

Commits on Sep 20, 2024

• Add support for `shard_axes` into `build_array_tspec_for_write`, to allow subchunking of arrays on specific axes when writing to storage.

  Show description for e37d633
  

  dicentra13

  authored and

  Orbax Authors

  committed
  e37d633

  Copy full SHA for e37d633

Commits on Sep 19, 2024

• Moves JaxModule.DEFAULT_METHOD_KEY into constants.py and updates references.

  Show description for fb0fc0b
  

  Orbax Authors

  committed
  fb0fc0b

  Copy full SHA for fb0fc0b

• Modernize type annotations using `from __future__ import annotations`.

  Show description for 1541e09
  

  niketkumar

  authored and

  Orbax Authors

  committed
  1541e09

  Copy full SHA for 1541e09

• No public description

  Show description for 228539c
  

  hawkinsp

  authored and

  Orbax Authors

  committed
  228539c

  Copy full SHA for 228539c

Commits on Sep 18, 2024

• Upgrade array write tspecs construction to use `build_array_tspec_for_write` in standard Orbax type handlers.

  Show description for f196924
  

  dicentra13

  authored and

  Orbax Authors

  committed
  f196924

  Copy full SHA for f196924

• Add a new helper function to construct TensorStore specs for writing arrays.

  Show description for 10eb878
  

  dicentra13

  authored and

  Orbax Authors

  committed
  10eb878

  Copy full SHA for 10eb878

Pagination

Previous

Next