Skip to content

Pull requests: jax-ml/jax

New pull request New
575 Open 22,898 Closed
575 Open 22,898 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

Update Jax B200 single gpu tests to use MIG partitions.
#33488 opened Nov 22, 2025 by copybara-service bot Loading…
Fix: Correct workflow path patterns and fix typo in CI workflows
#33486 opened Nov 22, 2025 by Anushavasa15 Loading…
5
Add portpicker to the freethreading deps and regenerate lock files. CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#33484 opened Nov 21, 2025 by emilyfertig Loading…
@emilyfertig
2
[test] explicitly specify gather mode in roofline test
#33483 opened Nov 21, 2025 by jakevdp Loading…
@jakevdp
3
Delete jax._src.dtypes.bit_width from internal API.
#33480 opened Nov 21, 2025 by copybara-service bot Loading…
@justinjfu
[Mosaic GPU] Support inference of additional layouts for BroadcastInDimOp.
#33473 opened Nov 21, 2025 by copybara-service bot Loading…
Reduce test_print_wgmma_tiled_layout shape size.
#33470 opened Nov 21, 2025 by copybara-service bot Loading…
[jax:docs] Update symbol for profiling.
#33463 opened Nov 21, 2025 by copybara-service bot Loading…
@danielsuo
[Pallas:MGPU] Fix plgpu.inline_mgpu support for scalar results.
#33462 opened Nov 21, 2025 by copybara-service bot Loading…
Automated Code Change
#33460 opened Nov 21, 2025 by copybara-service bot Loading…
Fix Ref indexing discrepancies with array indexing
#33456 opened Nov 21, 2025 by yashwantbezawada Loading…
1
10
Enable gather_with_index test on tc tiling. Also add a slicing to the test.
#33455 opened Nov 21, 2025 by copybara-service bot Loading…
@IvyZX
Reduce the maximum test parallelism on Windows machines to avoid OOMs
#33453 opened Nov 20, 2025 by copybara-service bot Loading…
@danielsuo
[CI] Testing changes to tpu jobs
#33447 opened Nov 20, 2025 by copybara-service bot Loading…
@MichaelHudgins
Add new BOUNDS_CHECK GatherScatterMode
#33446 opened Nov 20, 2025 by jakevdp Loading…
@jakevdp
4
Tiny usability tweak to allow user to pass a list instead of a tuple, e.g. pltpu.VMEM([8], jnp.int32).
#33439 opened Nov 20, 2025 by copybara-service bot Loading…
[Mosaic TPU] Remove logical-to-physical-device-id from tpu.td
#33438 opened Nov 20, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Remove pre-canonicalization optimization pass from tpu.td
#33437 opened Nov 20, 2025 by copybara-service bot Loading…
@apaszke
[Pallas/TPU] tpu_info.get_tpu_info optionally takes a TPU chip version argument
#33435 opened Nov 20, 2025 by copybara-service bot Loading…
Add metadata for CUDA and libtpu versions
#33431 opened Nov 20, 2025 by copybara-service bot Loading…
jax.typeof: improve documentation & add to API docs pull ready Ready for copybara import and testing
#33428 opened Nov 19, 2025 by jakevdp Loading…
@jakevdp
3
Seeing what happens if we get rid of positional effects pull ready Ready for copybara import and testing
#33425 opened Nov 19, 2025 by dougalm Loading…
3
[Mosaic] Add hasVectorOperandsOrResults utility function
#33422 opened Nov 19, 2025 by copybara-service bot Loading…
Check H100/B200 test results.
#33421 opened Nov 19, 2025 by copybara-service bot Loading…
@ybaturina
1
Change Tracer so it does not directly inherit from Array.
#33420 opened Nov 19, 2025 by copybara-service bot Loading…
@hawkinsp
ProTip! Mix and match filters to narrow down what you’re looking for.