Pull requests: jax-ml/jax
Pull requests list
[Pallas/Mosaic GPU][NFC] Slightly clarify the synchronization requirements in the docstring of
copy_gmem_to_smem. #33563 opened Nov 26, 2025 by copybara-service bot Loading…
Make sure axes are in the same order in pvary across multiple hosts. This ends up affecting
psum on the backward pass if axes order differs across hosts. #33562 opened Nov 26, 2025 by copybara-service bot Loading…
Enable test_major_scatter with TC tiling
#33554 opened Nov 26, 2025 by copybara-service bot Loading…
[export] Propagate the memory_space from out_shardings kokoro:force-run pull ready Ready for copybara import and testing
#33553 opened Nov 26, 2025 by gnecula Loading…
[xla:ffi] Move GPU FFI implementation to backends/gpu
#33549 opened Nov 26, 2025 by copybara-service bot Loading…
[Pallas TPU] Add mesh axis info to pallas call metadata
#33546 opened Nov 25, 2025 by copybara-service bot Loading…
[types] improve internal type annotations for indexing routines pull ready Ready for copybara import and testing
#33544 opened Nov 25, 2025 by jakevdp Loading…
Optimize lpmn_values using O(n²) Associated Legendre recurrences (#33370)
#33539 opened Nov 25, 2025 by AKash-A007 Loading…
Fixed docstring of make_array_from_process_local_data documentation
#33537 opened Nov 25, 2025 by vfdev-5 Loading…
[mosaic] tpu.memref_squeeze now verifies squeezed dimensions even when the source is not produced via tpu.erase_memref_layout
#33534 opened Nov 25, 2025 by copybara-service bot Loading…
[indexing] add
strategy argument to rewriting_take & implement static slicing strategy #33531 opened Nov 25, 2025 by jakevdp Loading…
Only consider Tracers with ShapedArray avals as instances of Array.
#33530 opened Nov 25, 2025 by copybara-service bot Loading…
ProTip! Type g i on any issue or pull request to go back to the issue listing page.