Pull requests: google/orbax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Replace deprecated
jax.tree_*
functions with jax.tree.*
#880
opened May 15, 2024 by
copybara-service
bot
Loading…
Incorporate async-compatible, process-subset-compatible barrier function. With this API, barrier names must be completely unique, and many places must be adjusted to take this into account. There is also an issue with inconsistency between
jax.process_index()
and distributed.client.process_id
, which is resolved by initialize_runtime_to_distributed_ids
. This function will be removed as soon as the discrepancy is resolved.
#879
opened May 14, 2024 by
copybara-service
bot
Loading…
Introduce
StepAlreadyExistsError
to be raised on save with existing step.
#878
opened May 14, 2024 by
copybara-service
bot
Loading…
Delegate to BasePyTreeCheckpointHandler rather than inheriting from it.
#876
opened May 10, 2024 by
copybara-service
bot
Loading…
Fork a small amount of Orbax code into Pax to deal with writing "aggregate" files, as Orbax will soon lose this ability.
#869
opened May 8, 2024 by
copybara-service
bot
Loading…
Add try catch block to catch ValueError reading from sharding string.
#857
opened May 1, 2024 by
copybara-service
bot
Loading…
Add OCDBT support to GlobalHostArray type_handler
#843
opened Apr 24, 2024 by
copybara-service
bot
Loading…
Create the utils function load_model_with_tfrt API for XLA devices.
#813
opened Apr 12, 2024 by
copybara-service
bot
Loading…
Enable Orbax checkpointing with background deleting thread
#796
opened Apr 4, 2024 by
copybara-service
bot
Loading…
Create an interface for retrieving the current step number.
#778
opened Mar 30, 2024 by
copybara-service
bot
Loading…
Changes default timeout in AsyncCheckpointer.
#698
opened Feb 1, 2024 by
copybara-service
bot
Loading…
Orbax test. Used for testing, don't submit.
#694
opened Jan 29, 2024 by
copybara-service
bot
Loading…
Fix errors of loading ckpt that contains sharding configs.
#671
opened Jan 17, 2024 by
copybara-service
bot
Loading…
Replace references to deprecated device_buffer attributes
#616
opened Dec 6, 2023 by
copybara-service
bot
Loading…
Previous Next
ProTip!
Mix and match filters to narrow down what you’re looking for.