Skip to content

Pull requests: google/orbax

New pull request New
101 Open 994 Closed
101 Open 994 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
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Moves TF dependent code from jax_module in to tensorflow_module and from jax_module_test to tensorflow_module_test. This code is duplicated for now and will be cleaned up in a subsequent cl.
#1196 opened Sep 19, 2024 by copybara-service bot Loading…
1
test restoring to global_mesh in parallel
#1195 opened Sep 19, 2024 by copybara-service bot Loading…
Add support for shard_axes into build_array_tspec_for_write, to allow subchunking of arrays on specific axes when writing to storage.
#1194 opened Sep 19, 2024 by copybara-service bot Loading…
Remove logging of ts_spec in validation step - this is very verbose.
#1191 opened Sep 18, 2024 by copybara-service bot Loading…
Internal Change Only
#1187 opened Sep 18, 2024 by copybara-service bot Loading…
1
Continues the TF dependency isolation refactor. The majority of the logic in JaxModule is wrapping jax in TF variables and functions. This logic is now moved in to the TensorFlowModule.
#1186 opened Sep 17, 2024 by copybara-service bot Loading…
1
Internal change.
#1183 opened Sep 17, 2024 by copybara-service bot Loading…
Use CheckpointHandlerRegistry when CheckpointManager is default-constructed, e.g. CheckpointManager(directory). Previously this usage pattern would default to single-item mode, but the new behavior will default to lazy initialization. The change should be a no-op for existing users.
#1177 opened Sep 13, 2024 by copybara-service bot Loading…
Use thread-safe barrier api by default in Orbax.
#1175 opened Sep 13, 2024 by copybara-service bot Loading…
Internal
#1150 opened Sep 9, 2024 by copybara-service bot Loading…
1
Internal change.
#1146 opened Sep 8, 2024 by copybara-service bot Loading…
1
Add an option to skip validating merged parameters in Orbax.
#1126 opened Sep 2, 2024 by copybara-service bot Loading…
Internal
#1125 opened Aug 31, 2024 by copybara-service bot Loading…
1
Internal change.
#1107 opened Aug 27, 2024 by copybara-service bot Loading…
Internal change.
#1095 opened Aug 23, 2024 by copybara-service bot Loading…
Internal.
#1090 opened Aug 22, 2024 by copybara-service bot Loading…
Move locking logic to snapshotting.
#1088 opened Aug 21, 2024 by copybara-service bot Loading…
Internal.
#1084 opened Aug 21, 2024 by copybara-service bot Loading…
internal changes only
#1079 opened Aug 20, 2024 by copybara-service bot Loading…
1
Remove deprecated _known_handlers
#1072 opened Aug 16, 2024 by copybara-service bot Loading…
1
Include directory path in create_directory barrier key.
#1040 opened Aug 5, 2024 by copybara-service bot Loading…
1
Add warning when GCS checkpoint is missing...
#1035 opened Jul 31, 2024 by rdyro Loading…
DO NOT SUBMIT
#1025 opened Jul 26, 2024 by copybara-service bot Loading…
1
Modernize typing in composite checkpoint handler.
#1014 opened Jul 22, 2024 by copybara-service bot Loading…
1
Sort devices explicitly by process index, then id (as opposed to IDs alone). IDs may be randomly generated, and are not guaranteed to be ordered based on their process index.
#996 opened Jul 10, 2024 by copybara-service bot Loading…
1
ProTip! Adding no:label will show everything without a label.