issues
search
google
/
orbax
Orbax provides common checkpointing and persistence utilities for JAX users
https://orbax.readthedocs.io/
Apache License 2.0
305
stars
36
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
Begins to wire up the new Orbax Model export path to allow exporting using the new Orbax Model pathway.
#1250
copybara-service[bot]
closed
1 month ago
1
Internal change.
#1249
copybara-service[bot]
closed
1 month ago
0
Orbax import breaks logging. Again.
#1248
qGentry
closed
1 month ago
7
Add `strict` option in `ArrayRestoreArgs`, defaulting to True. This prevents arrays from being accidentally padded or truncated when restoring.
#1247
copybara-service[bot]
closed
1 month ago
0
Internal change.
#1246
copybara-service[bot]
closed
1 month ago
0
Set `AsyncOptions.timeout_secs` default value to 10 minutes.
#1245
copybara-service[bot]
closed
1 month ago
0
GC potentially left over steps at startup time for local_ckpt_mgr.
#1244
copybara-service[bot]
closed
1 month ago
1
Reinstate subchunking logic to work around 0-sized dimensions.
#1243
copybara-service[bot]
closed
1 month ago
0
[NOT FOR MERGE] apply patch for _get_kvstore_for_yt
#1242
minotru
closed
1 month ago
1
Code duplication around tensorstore_spec logic in orbax-checkpoint
#1241
minotru
opened
1 month ago
2
Use `Fragments` for serialization.
#1240
copybara-service[bot]
closed
1 month ago
0
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.
#1239
copybara-service[bot]
closed
1 month ago
0
Move multihost package to _src/multihost. Important symbols are still exported publicly.
#1238
copybara-service[bot]
closed
1 month ago
0
Remove unused or test-only serialization functions.
#1237
copybara-service[bot]
closed
1 month ago
0
Rename `_single_item` to `_default_item` in `CheckpointManager`.
#1236
copybara-service[bot]
closed
1 month ago
0
Add a `SaveArgs` option that allows disabling pinned host transfer on a per-array basis.
#1235
copybara-service[bot]
closed
1 month ago
0
#cleanup Upgrade type annotations in `tensorstore_utils` to 3.10.
#1234
copybara-service[bot]
closed
1 month ago
0
Internal.
#1233
copybara-service[bot]
opened
1 month ago
0
Refactor calls to `choose_chunk_shape` to simplify `build_zarr_shard_and_chunk_metadata`.
#1232
copybara-service[bot]
closed
1 month ago
0
Internal change.
#1231
copybara-service[bot]
closed
1 month ago
0
Refactor `build_array_tspec_for_write` into a dataclass.
#1230
copybara-service[bot]
closed
1 month ago
0
Can't save model on Colab but can on Kaggle
#1229
windmaple
closed
1 month ago
1
Remove orbax-checkpoint support for Python 3.9.
#1228
copybara-service[bot]
closed
1 month ago
0
Furhter pushes TF dependencies down to the tensorflow_export. The serving_config processing now also occurs in the tensorflow_export class as it is strictly specific to TensorFlow. TF is still needed in the ExportManager since the public interface allows properties for getting the internal TF module and serving signatures.
#1227
copybara-service[bot]
closed
1 week ago
1
Fix `_determine_single_item_mode_from_directory` by incorporating step name format.
#1226
copybara-service[bot]
closed
1 month ago
0
Hi, I am trying to save checkpoints using the following code:
#1225
raresdolga
closed
1 month ago
3
Now that JAX supports donate on device_put, use that instead of the ugly device_put -> block -> delete routine
#1224
copybara-service[bot]
closed
1 month ago
0
Refactor reshard logic to a separate method for readability.
#1223
copybara-service[bot]
closed
1 month ago
1
Internal change.
#1222
copybara-service[bot]
opened
1 month ago
1
Fix slice selection logic in emergency checkpointing, so that a slice with a complete local checkpoint (if it exists) is always chosen as the secondary slice. Additionally, ensure that the `ArrayHandler` used to save persistent checkpoints is configured with the correct `primary_host`.
#1221
copybara-service[bot]
closed
1 month ago
0
Introduce `ListKey` and `TupleKey` for `list` and `tuple` PyTree nodes respectively.
#1220
copybara-service[bot]
closed
1 week ago
0
Rename CheckpointMetadata to StepMetadata.
#1219
copybara-service[bot]
closed
1 month ago
0
Restore from local with the mutated mesh, and transfer arrays to get back Pytree with original mesh before broadcasting.
#1218
copybara-service[bot]
closed
1 month ago
1
Refine several API docs and add a test to prevent colab docs from breaking.
#1217
copybara-service[bot]
closed
1 month ago
0
Support `empty tuple` values in checkpointable params.
#1216
copybara-service[bot]
closed
1 month ago
0
Support nested `asyncio.run` with `nest_asyncio` library.
#1215
copybara-service[bot]
closed
1 month ago
0
Make subchunking using `shard_axes` deterministic.
#1214
copybara-service[bot]
closed
1 month ago
0
Support nested `asyncio.run` without `nest_asyncio` library.
#1213
copybara-service[bot]
closed
1 month ago
0
Add monitoring to record how long directory creation takes (per sequential instance).
#1212
copybara-service[bot]
closed
1 month ago
0
Add `blocking_metadata_write` option to allow disabling async metadata write behavior.
#1211
copybara-service[bot]
closed
1 month ago
0
Restore from local with the mutated mesh, and transfer arrays to get back Pytree with original mesh before broadcasting.
#1210
copybara-service[bot]
closed
1 month ago
1
Use `__future__.annotations` module to modernize type definitions.
#1209
copybara-service[bot]
closed
2 months ago
0
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 slightly reduced verbosity if we just want to ensure an association between an item name and args/handler.
#1208
copybara-service[bot]
closed
1 month ago
0
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 an error if a requested item does not exist, and will attempt to restore all existing items if empty `args` are provided.
#1207
copybara-service[bot]
closed
2 months ago
0
No public description
#1206
copybara-service[bot]
opened
2 months ago
1
Fix a few typos in the comments.
#1205
copybara-service[bot]
closed
2 months ago
1
Adds NotImplemented errors to obm_module and adds related tests to jax_module_test.
#1204
copybara-service[bot]
closed
2 months ago
1
Refactor to extract a separate module, `asyncio_utils`, for asyncio helper functions from `path/async_utils` module.
#1203
copybara-service[bot]
closed
2 months ago
0
Log error if number of checkpoints exceeds max_to_keep instead of raising error in emergency checkpointing.
#1202
copybara-service[bot]
closed
2 months ago
0
Bump ipython from 7.23.1 to 8.10.0 in /docs/requirements in the pip group across 1 directory
#1201
dependabot[bot]
opened
2 months ago
0
Previous
Next