Skip to content

Pull requests: google/orbax

New pull request New
204 Open 1,804 Closed
204 Open 1,804 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Use gigabytes for jax checkpoint write/read metrics.
#2150 opened Jul 22, 2025 by copybara-service bot Loading…
Copy Bf16Optimizaion implementation to Orbax.
#2149 opened Jul 22, 2025 by copybara-service bot Loading…
Export multiple stablehlos for a JAX poly function, based on discrete values for symbols provided by user
#2148 opened Jul 22, 2025 by copybara-service bot Loading…
1
Fix load_pytree() incompoatibility bug.
#2147 opened Jul 21, 2025 by copybara-service bot Loading…
Small refactoring: Pass exactly what the function needs
#2146 opened Jul 21, 2025 by copybara-service bot Loading…
1
Remove the unused TemporaryPath.match method.
#2143 opened Jul 21, 2025 by copybara-service bot Loading…
Internal change
#2141 opened Jul 21, 2025 by copybara-service bot Loading…
1
Move blocking save implementation details from saving to saving_utils.
#2140 opened Jul 18, 2025 by copybara-service bot Loading…
v1: Update LeafHandlerRegistry to work with handler_type only.
#2133 opened Jul 18, 2025 by copybara-service bot Loading…
Move implementation details of Orbax V1 save_pytree and save_checkpointables into helper file.
#2127 opened Jul 17, 2025 by copybara-service bot Loading…
Fixes for safetensors import error from previous rollback.
#2125 opened Jul 17, 2025 by copybara-service bot Loading…
Add a new parameter background_thread_count to allow users to set the number of threads for ThreadedCheckpointDeleter.
#2116 opened Jul 16, 2025 by copybara-service bot Loading…
[REX] Enable ttl based deletion of step and substeps
#2114 opened Jul 16, 2025 by copybara-service bot Loading…
[Orbax Model] Save XLA compile options per platform.
#2113 opened Jul 15, 2025 by copybara-service bot Loading…
Add an unit test for oex_orchestration.create().
#2108 opened Jul 14, 2025 by copybara-service bot Loading…
1
Use Checkpoint Manager to benchmark in Orbax V0.
#2102 opened Jul 14, 2025 by copybara-service bot Loading…
Internal change
#2101 opened Jul 12, 2025 by copybara-service bot Loading…
1
Use HNS-compatible listing in NameFormat.
#2098 opened Jul 11, 2025 by copybara-service bot Loading…
Attempt to restore from single replica instead of single slice
#2096 opened Jul 11, 2025 by findmyway • Draft
Add branching logic to loading.py and refactor the checkpoint handler logic.
#2094 opened Jul 10, 2025 by copybara-service bot Loading…
Swap order of expected value and actual value when using [assertProto2Equal](https://source.corp.google.com/piper///depot/google3/net/proto2/contrib/pyutil/compare.py;rcl=781699748;l=153) to be consistent with definition.
#2092 opened Jul 10, 2025 by copybara-service bot Loading…
1
Make the temp directory generation be unique across CheckpointManager when shared barrier_sync_key_prefix.
#2086 opened Jul 9, 2025 by copybara-service bot Loading…
1
Restore from single slice
#2084 opened Jul 9, 2025 by findmyway • Draft
9
Add CheckpointLayout protocol for various checkpoint formats to inherit from.
#2081 opened Jul 7, 2025 by copybara-service bot Loading…
Pr 2057
#2074 opened Jul 2, 2025 by deepikarajani24 Loading…
ProTip! Mix and match filters to narrow down what you’re looking for.