Skip to content

Pull requests: NVIDIA/JAX-Toolbox

New pull request New
41 Open 1,828 Closed
41 Open 1,828 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

Read STAGING.md file from xla/nv-staging/latest and define JAX commit
#2054 opened Apr 24, 2026 by Steboss Contributor Loading…
Nvidia-runners S3 cache implementation
#2053 opened Apr 23, 2026 by Steboss Contributor Loading…
1
nsys-jax: fix collective metadata extraction.
#2048 opened Apr 21, 2026 by olupton Collaborator Loading…
1
Add EKS JobSet action
#2037 opened Apr 15, 2026 by aybchan Member Loading…
2 tasks done
@aybchan
3
Add cutile as part of JAX dependency
#2020 opened Apr 1, 2026 by gpupuck Contributor Loading…
Add GKE cluster check in composite action
#2017 opened Mar 30, 2026 by aybchan Member Loading…
@aybchan
Upgrade base CUDA to 13.2
#2015 opened Mar 25, 2026 by aybchan Member Loading…
1
Refactor of vLLM bridge
#1914 opened Jan 23, 2026 by keshavb96 Contributor Loading…
7
Install uv (https://docs.astral.sh/uv/) in base container
#1891 opened Jan 12, 2026 by andportnoy Contributor Loading…
Upgrade to xpk v1.0.0, disable nccl installation on host on cluster creation
#1881 opened Jan 7, 2026 by aybchan Member Loading…
@aybchan
2
[jax-inference-offloading] consolidate definitions for default tensor dtype
#1816 opened Dec 5, 2025 by yhtang Contributor Loading…
8
[jax-inference-offloading] Bump to CUDA 13.0
#1807 opened Dec 2, 2025 by yhtang Contributor • Draft
Add basic sitrep steps for GKE XPK action
#1579 opened Jul 25, 2025 by aybchan Member Loading…
@aybchan
Add MACE training example
#1192 opened Dec 10, 2024 by mariogeiger Loading…
Customize NCCL for base container
#1123 opened Oct 28, 2024 by DwarKapex Contributor Loading…
2
ci[maxtext]: adds rosetta directory to maxtext containers
#1097 opened Oct 13, 2024 by terrykong Contributor Loading…
2
chore: remove deprecated flags from model xla flag configs
#1088 opened Oct 8, 2024 by terrykong Contributor Loading…
1
Fix the Order of GatedAct to be (act, linear) when Pax2TE.
#892 opened Jun 10, 2024 by mingxu1067 Contributor Loading…
Add upstream pax 2-node Grok test
#854 opened May 24, 2024 by ashors1 Contributor Loading…
2
Support auto-generated FP8 meta for CKPT converters.
#844 opened May 21, 2024 by mingxu1067 Contributor Loading…
4
Add FMHA PAXML test
#830 opened May 16, 2024 by hmonishN Contributor Loading…
33
Enable TE fused attention by default in rosetta pax container
#757 opened Apr 23, 2024 by ashors1 Contributor Loading…
Experiment with IB/OFED libs from distro repo
#756 opened Apr 23, 2024 by yhtang Contributor Loading…
ci: Restore rosetta-t5x-test
#749 opened Apr 22, 2024 by ko3n1g Contributor Loading…
4
feat: Add rosetta-maxtext
#738 opened Apr 18, 2024 by ko3n1g Contributor Loading…
16
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.