Skip to content

Pull requests: jax-ml/jax

New pull request New
417 Open 16,993 Closed
417 Open 16,993 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

[better_errors] Continue adding debug info to Jaxprs (step 3) pull ready Ready for copybara import and testing
#26348 opened Feb 6, 2025 by gnecula Loading…
@gnecula
Add type annotation for zero_from_primal
#26346 opened Feb 6, 2025 by NeilGirdhar Loading…
Scaled matmul for mxfp8
#26345 opened Feb 5, 2025 by wenscarl • Draft
disable cuDNN flash attention head dim 256 on blackwell
#26344 opened Feb 5, 2025 by Cjkkkk Loading…
1
doc: improve docs for jax.lax trig functions documentation
#26342 opened Feb 5, 2025 by jakevdp Loading…
@jakevdp
Port Csrlsvqr GPU kernel to the FFI.
#26341 opened Feb 5, 2025 by copybara-service bot Loading…
@dfm
[pallas:mosaic_gpu] Fixed a bug in how delay_release is handled in emit_pipeline
#26337 opened Feb 5, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Add tests for various tcgen05.mma configurations
#26335 opened Feb 5, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add more SMEM buffers to avoid blocking for memory traffic
#26334 opened Feb 5, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Take TMEM as a TMEMRef in tcgen05.mma, not as a raw address
#26333 opened Feb 5, 2025 by copybara-service bot Loading…
@apaszke
Check "jax_rocm_visible_devices" at client creation. pull ready Ready for copybara import and testing
#26327 opened Feb 5, 2025 by ksebaz Loading…
@dfm
1
[JAX] Generate more readable error for failed device deserialization in colocated Python
#26321 opened Feb 4, 2025 by copybara-service bot Loading…
Test ci changes
#26314 opened Feb 4, 2025 by copybara-service bot Loading…
Finalize some deprecations from jax.lib.xla_client
#26292 opened Feb 4, 2025 by jakevdp • Draft
@jakevdp
3
Extend random.orthogonal to semi-orthogonal matrices and simplify initializers.orthogonal pull ready Ready for copybara import and testing
#26291 opened Feb 4, 2025 by carlosgmartin Loading…
@jakevdp
8
[JAX][DOC] memory_kind, with_memory_kind and out_shardings pull ready Ready for copybara import and testing
#26285 opened Feb 3, 2025 by zhenying-liu Loading…
@yashk2810
2
Proof-of-concept: Use XLA's intra_op_threadpool to batch parallelize JAX's LAPACK calls on CPU.
#26280 opened Feb 3, 2025 by copybara-service bot Loading…
@dfm
[CI] Add upload to resultstore to tsan bazel test run CI
#26279 opened Feb 3, 2025 by MichaelHudgins Loading…
1
[pallas:mgpu] Change FA3 kernel bc lax.div doesn't like mixed types anymore.
#26268 opened Feb 3, 2025 by copybara-service bot Loading…
Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'.
#26264 opened Feb 2, 2025 by copybara-service bot Loading…
@skye
Add check to lax.composite to prevent DynamicJaxprTracer type errors.
#26257 opened Feb 1, 2025 by copybara-service bot Loading…
Add num_updates dimension to input/output_offset and send/recv_sizes of ragged_all_to_all.
#26256 opened Feb 1, 2025 by copybara-service bot Loading…
Return arrays from ArrayImpl._check_and_rearrange. Build IFRT shardings with both addressable and non-addressable devices, instead of only addressable devices.
#26253 opened Jan 31, 2025 by copybara-service bot Loading…
@emilyfertig
Make jnp.asarray lower to asarray_p in the simplest cases.
#26244 opened Jan 31, 2025 by jakevdp • Draft
@jakevdp
[mosaic_gpu/pallas:mgpu] Eradicate wgmma_layout
#26240 opened Jan 31, 2025 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.