Skip to content

Pull requests: jax-ml/jax

New pull request New
497 Open 19,009 Closed
497 Open 19,009 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
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[XLA:GPU] Do not do dot strength reduction for vector×matrix dots.
#28882 opened May 21, 2025 by copybara-service bot Loading…
Introduce profiler_options in the documentation.
#28880 opened May 21, 2025 by copybara-service bot Loading…
[Pallas:MGPU] Disable the GPU distributed test when the platform allocator is used CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28879 opened May 21, 2025 by copybara-service bot Loading…
@apaszke
Expose get_mesh_shape API for HloSharding such that we can call it directly in Python with JAX.
#28875 opened May 21, 2025 by copybara-service bot Loading…
add jax cudnn sdpa mla support
#28872 opened May 21, 2025 by Cjkkkk Loading…
Introduce the flag add_pypi_cuda_wheel_deps that controls if the tests depend on NVIDIA CUDA wheels hermetically.
#28870 opened May 21, 2025 by copybara-service bot Loading…
[Mosaic:TPU][Relayout] Remove minor implicit dimension for 32-bit native tiling
#28868 opened May 21, 2025 by copybara-service bot Loading…
Prototype of cross-host device transfers in IFRT-PJRT. For now it only works with the TFRT TPU runtime, because other PjRt plugins don't implement the necessary APIs.
#28867 opened May 20, 2025 by copybara-service bot Loading…
@emilyfertig
Adjust triton dialect lowering rounding mode to allow upcasting fp8 types
#28866 opened May 20, 2025 by copybara-service bot Loading…
1
@rdyro
Fix native tiling logic in infer_vector_layout.
#28862 opened May 20, 2025 by copybara-service bot Loading…
Create workflow for testing bazel cuda non rbe + py import
#28860 opened May 20, 2025 by copybara-service bot Loading…
2
[Pallas/Mosaic GPU] Add a Pallas:MGPU implementation of ragged_dot.
#28853 opened May 20, 2025 by copybara-service bot Loading…
@bchetioui
[pallas] Slightly revamped how checkify is exposed in Pallas
#28852 opened May 20, 2025 by copybara-service bot Loading…
@superbobry
[Pallas:MGPU] Make semaphores compatible with plgpu.kernel and the profiler
#28847 opened May 20, 2025 by copybara-service bot Loading…
@apaszke
Experiment: mutable RNG object pull ready Ready for copybara import and testing
#28845 opened May 20, 2025 by jakevdp • Draft
@jakevdp
[ragged-paged-attn] Use select for initialization in flash attention.
#28840 opened May 19, 2025 by copybara-service bot Loading…
In print_environment_info, print environment variables that start with XLA_
#28839 opened May 19, 2025 by carlosgmartin Loading…
Prevent _compute_fans from raising a ZeroDivisionError if a dimension of shape is 0.
#28837 opened May 19, 2025 by copybara-service bot Loading…
Bump the pip group across 1 directory with 2 updates dependencies Pull requests that update a dependency file python Pull requests that update Python code
#28835 opened May 19, 2025 by dependabot bot Loading…
Bump hypothesis from 6.102.4 to 6.131.18 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#28831 opened May 19, 2025 by dependabot bot Loading…
@jakevdp
#sdy Always export Shardy when going to HLO.
#28821 opened May 19, 2025 by copybara-service bot Loading…
[pallas] Fix broadcast_in_dim fuser eval rule.
#28819 opened May 19, 2025 by copybara-service bot Loading…
@chr1sj0nes
added solve_sylvester and accompanying tests
#28810 opened May 17, 2025 by Dekermanjian Loading…
@jakevdp
17
[Mosaic GPU] Add a "realistic" blackwell matmul example. Also implements the unimplemented ops required to make this work:
#28806 opened May 17, 2025 by copybara-service bot Loading…
@justinjfu
Enable the iota tests in pallas, remove broken 1d impl
#28801 opened May 16, 2025 by copybara-service bot Loading…
ProTip! Mix and match filters to narrow down what you’re looking for.