google orbax issues - Githubissues

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users

https://orbax.readthedocs.io/

Apache License 2.0

304 stars 36 forks source link

issues

Newest

Newest Most commented Recently updated Oldest Least commented Least recently updated

Weights of Restored model after checkpointing do not give same loss as model before saving

#1402 delara38 opened 8 hours ago
0
internal change.

#1401 copybara-service[bot] opened 8 hours ago
0
internal change.

#1400 copybara-service[bot] opened 10 hours ago
0
Fix `tf.Module` tracking problem.

#1399 copybara-service[bot] closed 7 hours ago
0
Update CompositeCheckpointHandler.metadata() to return StepMetadata.

#1398 copybara-service[bot] opened 15 hours ago
0
Internal change.

#1397 copybara-service[bot] closed 15 hours ago
0
fix requirements for jax.

#1396 copybara-service[bot] opened 16 hours ago
0
Fix class name of a namedtuple created on the fly by replacing `.` with `_`.

#1395 copybara-service[bot] closed 16 hours ago
0
Fix metadata ser/deser of jax registered container nodes like flax.struct.

#1394 copybara-service[bot] opened 17 hours ago
0
Internal

#1393 copybara-service[bot] opened 18 hours ago
0
Fix local restore by re-mapping device ids directly instead of inferring them from how process indexes changed across restarts with some false assumptions.

#1392 copybara-service[bot] opened 19 hours ago
0
Minor fixes to metadata serialization and updates.

#1391 copybara-service[bot] opened 19 hours ago
0
Internal

#1390 copybara-service[bot] opened 19 hours ago
0
Reflect https://github.com/google/orbax/pull/1358, which enables replica_parallel saving. Uses cache to avoid recomputing `devices_indices_map` for arrays with identical shape/sharding.

#1389 copybara-service[bot] closed 17 hours ago
0
Fix namedtuple empty value typestr when experimental support_rich_types is disabled again after enabling it.

#1388 copybara-service[bot] closed 20 hours ago
0
Rollback Layout support in StandardCheckpointHandler

#1387 copybara-service[bot] closed 1 day ago
0
Add support to Layout using StandardCheckpointHandler

#1386 copybara-service[bot] closed 23 hours ago
0
Internal change.

#1385 copybara-service[bot] closed 7 hours ago
0
fix BUILD files.

#1384 copybara-service[bot] closed 3 days ago
0
Removes the build file out of the obm_module and instead processes the serving_configs in the constructor of obm_export to match what is done with the Tensorflow Export.

#1383 copybara-service[bot] opened 3 days ago
1
Add docstring for CheckpointManagerOptions.temporary_path_class.

#1382 copybara-service[bot] closed 3 days ago
1
rename workspace.BUILD to BUILD.

#1381 copybara-service[bot] closed 3 days ago
0
Add tests and documentation for abstract_arrays.

#1380 copybara-service[bot] closed 3 days ago
0
Internal change.

#1379 copybara-service[bot] closed 3 days ago
1
Release 0.10.1.

#1378 copybara-service[bot] closed 3 days ago
0
Increase version to 0.10.1 to incorporate emergency checkpointing fix.

#1377 copybara-service[bot] closed 3 days ago
0
Correct issue where emergency checkpoint debug logging was not OSS-ed.

#1376 copybara-service[bot] closed 3 days ago
0
Add /checkpoint/_src/path/snapshot/BUILD and /checkpoint/_src/path/BUILD.

#1375 copybara-service[bot] closed 3 days ago
0
Add support to Layout using StandardCheckpointHandler

#1374 copybara-service[bot] closed 3 days ago
0
Release 0.10.0 and correct the changelog.

#1373 copybara-service[bot] closed 4 days 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 3 days ago
0
`ServingConfig.extra_trackable_resources` is now supported in the Orbax Model path when `support_tf_resources` is `True`.

#1371 copybara-service[bot] closed 3 days ago
0
Add logging to detect missing chunks to emergency checkpointing to facilitate local checkpoint debugging.

#1370 copybara-service[bot] closed 4 days ago
0
Need to add because parent CL updates Orbax codebase.

#1369 copybara-service[bot] opened 4 days ago
0
Add an option to disable saving root-level metadata in CheckpointManager.

#1368 copybara-service[bot] closed 3 days ago
0
Change XlaRuntimeError to jax.errors.JaxRuntimeError

#1367 copybara-service[bot] closed 4 days ago
0
setup github actions to periodicly run existing OSS tests every hour.

#1366 copybara-service[bot] closed 4 days ago
0
Internal

#1365 copybara-service[bot] opened 5 days 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 5 days ago
0
Add .bazelrc, .bazelversion, requirements.txt, WORKSPACE to orbax github repo.

#1363 copybara-service[bot] closed 5 days ago
0
[Experimental Feature] Support `empty NamedTuple` leaf when `PyTreeMetadataOptions.support_rich_types=true`.

#1362 copybara-service[bot] closed 3 days ago
0
Test, this shoudl fail.

#1361 copybara-service[bot] closed 3 days ago
0
Expose public symbols via orbax/checkpoint/type_handlers.py

#1360 copybara-service[bot] closed 4 days ago
0
Change XlaRuntimeError to jax.errors.JaxRuntimeError

#1359 jakevdp closed 4 days ago
1
[replica-parallel] Avoid jax.sharding.Sharding#devices_indices_map

#1358 gspschmid closed 17 hours ago
3
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 5 days ago
0
`ocp.tree.serialize_tree` filtering logic for sequences with empty leaves

#1356 JesseFarebro opened 1 week ago
1
Add/Update tests for `is_empty_or_leaf` and `is_empty_node`.

#1355 copybara-service[bot] closed 6 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 5 days ago
0
Create a separate namespace package `orbax.checkpoint.testing` for exporting test objects.

#1353 copybara-service[bot] closed 1 week ago
0