Skip to content

Pull requests: apple/axlearn

New pull request New
49 Open 1,117 Closed
49 Open 1,117 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

Support launching CPU jobset
#1223 opened May 29, 2025 by muyangyuapple Loading…
Logmel: use jnp.fft.rfft instead of jnp.fft.fft.
#1221 opened May 28, 2025 by ds-hwang Loading…
1
Ensure GCS bucket region matches config region
#1220 opened May 28, 2025 by ehorning Loading…
Add max_len arg to fake_speech_source().
#1216 opened May 27, 2025 by ds-hwang Loading…
1
Add DataflowRunnerJob to allow running Dataflow jobs via Bastion.
#1213 opened May 27, 2025 by zachberger • Draft
4 tasks
[JAX API Update] Remove jax_spmd_mode from config
#1211 opened May 23, 2025 by Steboss Loading…
Fuji with Orbax Emergency Checkpointer
#1210 opened May 23, 2025 by samos123 • Draft
Updated jax.tree_map to jax.tree_util.tree_map in update_transformation.py
#1207 opened May 22, 2025 by m-ny Loading…
[UPDATE JAX API] Modify all the calls to jax.tree
#1206 opened May 22, 2025 by Steboss Loading…
[Do not review] Fuji inference examples
#1204 opened May 21, 2025 by hanzhi713 • Draft
[UPDATE JAX API] Update trainer_test _normalized_spec
#1201 opened May 21, 2025 by Steboss Loading…
2
[UPDATE JAX API] Update jax.experimental.array_serialization.serialization to latest JAX
#1198 opened May 21, 2025 by Steboss Loading…
1
Debug regression
#1196 opened May 21, 2025 by raytomato • Draft
3
Add logging for dataset input and checkpoint saving
#1190 opened May 18, 2025 by linxiulei Loading…
13
Dynamically generate FP8 configs for Fuji
#1186 opened May 15, 2025 by andersensam Loading…
3
Jax 0.5.3
#1182 opened May 14, 2025 by matthew-e-hopkins Loading…
Adding LWS Integration
#1174 opened May 12, 2025 by Edwinhr716 Loading…
19
Update Dockerfile to use ubuntu:22.04
#1169 opened May 8, 2025 by andersensam Loading…
9
Use lambda instead of functools.partial to create single-host config maker
#1152 opened May 2, 2025 by yhtang Loading…
1
7
Save O_PROJ on Fuji 70B-v2 for TRN2
#1147 opened May 1, 2025 by apoorvtintin Loading…
4
Fallback to standard mesh on neuron backend for incompatible multi-granule meshes
#1146 opened May 1, 2025 by apoorvtintin Loading…
5
repro PR for test_spmd_evaler fail with jax 0.5.3
#1114 opened Apr 18, 2025 by matthew-e-hopkins • Draft
Add ROCm support
#1112 opened Apr 16, 2025 by akasharidas Loading…
14
Migrate from Legacy JAX APIs jax.tree_util to jax.tree
#986 opened Feb 12, 2025 by apivovarov Loading…
7
Adds logging to aux loss collection.
#950 opened Jan 26, 2025 by ruomingp Loading…
1
ProTip! Exclude everything labeled bug with -label:bug.