issues
search
google
/
orbax
Orbax provides common checkpointing and persistence utilities for JAX users
https://orbax.readthedocs.io/
Apache License 2.0
303
stars
36
forks
source link
issues
Newest
Newest
Most commented
Recently updated
Oldest
Least commented
Least recently updated
Add support to Layout using StandardCheckpointHandler
#1374
copybara-service[bot]
closed
9 minutes ago
0
Release 0.10.0 and correct the changelog.
#1373
copybara-service[bot]
closed
12 hours ago
0
Add `shape` property to `ArrayRestoreArgs` to be interchangeable with `global_shape`. This allows generic functions like `to_shape_dtype_struct` to reference it.
#1372
copybara-service[bot]
closed
5 minutes ago
0
`ServingConfig.extra_trackable_resources` is now supported in the Orbax Model path when `support_tf_resources` is `True`.
#1371
copybara-service[bot]
opened
16 hours ago
0
Add logging to detect missing chunks to emergency checkpointing to facilitate local checkpoint debugging.
#1370
copybara-service[bot]
closed
16 hours ago
0
Need to add because parent CL updates Orbax codebase.
#1369
copybara-service[bot]
opened
19 hours ago
0
Add an option to disable saving root-level metadata in CheckpointManager.
#1368
copybara-service[bot]
opened
19 hours ago
0
Change XlaRuntimeError to jax.errors.JaxRuntimeError
#1367
copybara-service[bot]
closed
22 hours ago
0
setup github actions to periodicly run existing OSS tests every hour.
#1366
copybara-service[bot]
closed
21 hours ago
0
Internal
#1365
copybara-service[bot]
opened
1 day ago
0
Add vlogs for emergency checkpointing so we can easily tell which indices of the array need to be loaded.
#1364
copybara-service[bot]
closed
1 day ago
0
Add .bazelrc, .bazelversion, requirements.txt, WORKSPACE to orbax github repo.
#1363
copybara-service[bot]
closed
1 day ago
0
[Experimental Feature] Support `empty NamedTuple` leaf when `PyTreeMetadataOptions.support_rich_types=true`.
#1362
copybara-service[bot]
opened
1 day ago
0
Test, this shoudl fail.
#1361
copybara-service[bot]
opened
1 day ago
0
Expose public symbols via orbax/checkpoint/type_handlers.py
#1360
copybara-service[bot]
closed
22 hours ago
0
Change XlaRuntimeError to jax.errors.JaxRuntimeError
#1359
jakevdp
closed
22 hours ago
1
[replica-parallel] Avoid jax.sharding.Sharding#devices_indices_map
#1358
gspschmid
opened
1 day ago
2
Merge https://github.com/google/orbax/pull/1320 which adds replica-parallel saving for arrays with replicated shards. Thanks to https://github.com/gspschmid.
#1357
copybara-service[bot]
closed
1 day ago
0
`ocp.tree.serialize_tree` filtering logic for sequences with empty leaves
#1356
JesseFarebro
opened
3 days ago
1
Add/Update tests for `is_empty_or_leaf` and `is_empty_node`.
#1355
copybara-service[bot]
closed
2 days ago
0
Adds a new `support_tf_resources` flag (defaults to `False`) to `jax2obm_kwargs` to enable a new way to save TF (concrete) functions in Orbax Model that allows those functions to use TF resources (e.g. `tf.Variable`, vocabulary tables, embedding tables, etc.).
#1354
copybara-service[bot]
closed
1 day ago
0
Create a separate namespace package `orbax.checkpoint.testing` for exporting test objects.
#1353
copybara-service[bot]
closed
3 days ago
0
Internal
#1352
copybara-service[bot]
closed
2 days ago
1
Internal change.
#1351
copybara-service[bot]
closed
3 days ago
0
Internal
#1350
copybara-service[bot]
opened
3 days ago
0
Add support for legacy metadata file path in CheckpointManager.
#1349
copybara-service[bot]
closed
3 days ago
0
Internal change.
#1348
copybara-service[bot]
closed
3 days ago
0
Consider NamedTuples as concrete class and never empty.
#1347
copybara-service[bot]
closed
1 day ago
0
Refactor and restructure constructs from `type_handlers.py` and `metadata`
#1346
copybara-service[bot]
closed
1 day ago
0
[Experimental Feature] Support `NamedTuple` and `Tuple` nodes in PyTree metadata.
#1345
copybara-service[bot]
closed
1 day ago
0
Wrap `tree.utils.serialize_tree` -> `metadata.tree.serialize_tree` to allow `NamedTuple` and `Tuple` node types in PyTree metadata.
#1344
copybara-service[bot]
closed
3 days ago
0
Move `Checkpointer` implementations to `_src`.
#1343
copybara-service[bot]
closed
2 days ago
0
Checking changes back in with a bug fix.
#1341
copybara-service[bot]
closed
6 days ago
1
Add chex as an optional testing dependency to unbreak RTD. Unclear why this has not previously been a problem.
#1340
copybara-service[bot]
closed
6 days ago
0
[emergency checkpoint] Add `ReplicatorCheckpointManager` implementation for interoperating with replicator service provided by GKE (or theoretically, any other similar service).
#1339
copybara-service[bot]
opened
6 days ago
0
Internal change.
#1338
copybara-service[bot]
closed
6 days ago
0
Orbax Checkpointing and Flax.NNX require hacking to work together
#1337
hdrwilkinson
opened
1 week ago
1
Rolling back previous changes due to a breakage.
#1336
copybara-service[bot]
closed
1 week ago
1
Add user-facing Root/StepMetadata classes and internal serialization modules.
#1335
copybara-service[bot]
closed
19 hours ago
0
Updated Version
#1334
copybara-service[bot]
closed
1 week ago
1
Remove locking and its dependencies.
#1333
copybara-service[bot]
closed
1 week ago
0
Break out mesh construction and process ID metadata utils into a separate file.
#1331
copybara-service[bot]
closed
1 week ago
0
Submission of https://github.com/google/orbax/pull/1319. All credit and thanks to https://github.com/gspschmid.
#1330
copybara-service[bot]
closed
1 week ago
0
Add Layout support
#1329
copybara-service[bot]
closed
1 week ago
0
Add MultiProcessTest base.
#1328
copybara-service[bot]
opened
1 week ago
0
Fix device to HBM mapping for Trillium.
#1327
copybara-service[bot]
closed
1 week ago
0
Custom TypeHandler and "No per-process OCDBT checkpoint subdirs" warning
#1326
PhilipVinc
opened
1 week ago
2
Internal change.
#1325
copybara-service[bot]
closed
1 week ago
0
Manual submit of https://github.com/google/orbax/pull/1323.
#1324
copybara-service[bot]
closed
1 week ago
0
Resubmit single array error message
#1323
cpgaffney1
opened
1 week ago
0
Next