Skip to content

Pull requests: NVIDIA/JAX-Toolbox

New pull requestNew
42 Open 1,614 Closed
42 Open 1,614 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 reviewsReview requiredApproved reviewChanges requested
Assignee
Filter by who’s assigned
Assigned to nobodyLoading
Sort
Sort by
NewestOldestMost commentedLeast commentedRecently updatedLeast recently updatedBest match
Most reactions
👍👎😄🎉😕❤️🚀👀

Pull requests list

Port over nvcr login and release flag changes
#1840 opened Dec 10, 2025 by arjkeshLoading…
cuda-dl-base: bump to 25.11
#1837 opened Dec 10, 2025 by oluptonLoading…
Nsys JAX update for remapping nccl nodes
#1829 opened Dec 9, 2025 by StebossLoading…
1
[jax-inference-offloading] Update vllm version to 0.12.0
#1825 opened Dec 8, 2025 by yhtangLoading…
[jax-inference-offloading] consolidate definitions for default tensor dtype
#1816 opened Dec 5, 2025 by yhtangLoading…
8
[jax-inference-offloading] Bump to CUDA 13.0
#1807 opened Dec 2, 2025 by yhtang • Draft
JAX-vLLM Offloading k8s (GKE)
#1797 opened Nov 24, 2025 by aybchan • Draft
@aybchan
13
Always use KMD over UMD for JAX tests
#1794 opened Nov 24, 2025 by gpupuckLoading…
2
Build and test for AXLearn on ARM64
#1763 opened Nov 3, 2025 by StebossLoading…
25.10 NGC release testing
#1693 opened Sep 26, 2025 by aybchanLoading…
1
Test NCCL with EKS
#1651 opened Sep 1, 2025 by Steboss • Draft
Add basic sitrep steps for GKE XPK action
#1579 opened Jul 25, 2025 by aybchanLoading…
@aybchan
25.08: Add NGC release testing
#1554 opened Jul 15, 2025 by aybchanLoading…
Add MACE training example
#1192 opened Dec 10, 2024 by mariogeigerLoading…
Customize NCCL for base container
#1123 opened Oct 28, 2024 by DwarKapexLoading…
2
ci[maxtext]: adds rosetta directory to maxtext containers
#1097 opened Oct 13, 2024 by terrykongLoading…
2
Add cuda profiler hook to paxml
#1089 opened Oct 9, 2024 by shawnwang18Loading…
6
chore: remove deprecated flags from model xla flag configs
#1088 opened Oct 8, 2024 by terrykongLoading…
1
Fix the Order of GatedAct to be (act, linear) when Pax2TE.
#892 opened Jun 10, 2024 by mingxu1067Loading…
Add upstream pax 2-node Grok test
#854 opened May 24, 2024 by ashors1Loading…
2
Support auto-generated FP8 meta for CKPT converters.
#844 opened May 21, 2024 by mingxu1067Loading…
4
Add FMHA PAXML test
#830 opened May 16, 2024 by hmonishNLoading…
33
Enable TE fused attention by default in rosetta pax container
#757 opened Apr 23, 2024 by ashors1Loading…
Experiment with IB/OFED libs from distro repo
#756 opened Apr 23, 2024 by yhtangLoading…
ci: Restore rosetta-t5x-test
#749 opened Apr 22, 2024 by ko3n1gLoading…
4
ProTip! Updated in the last three days: updated:>2025-12-10.