Skip to content

Pull requests: jax-ml/jax

New pull request New
539 Open 22,495 Closed
539 Open 22,495 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

nn: use ensure_arraylike instead of check_arraylike and jnp.asarray
#32990 opened Oct 30, 2025 by arvindsankar Loading…
3
[Mosaic TPU] Support more general i1 relayout which usually comes from masking. Before this change, we will pick too small sublane tiling for packed types during relayout insertion, e.g., choose (1, 128) for i8.
#32989 opened Oct 30, 2025 by copybara-service bot Loading…
[Mosaic TPU] Add pack/unpack elementwise operations.
#32988 opened Oct 29, 2025 by copybara-service bot Loading…
Create an internal LowerType that can be passed to _resolve_and_lower via Traced and normal jit non-AOT path instead of actual arguments. We still hold onto consts though in Traced.
#32987 opened Oct 29, 2025 by copybara-service bot Loading…
@yashk2810
Added transpose rhs option to ragged dot kernel, useful for implementing the backward function.
#32986 opened Oct 29, 2025 by copybara-service bot Loading…
Test hermetic C++ container
#32985 opened Oct 29, 2025 by copybara-service bot Loading…
@quoctruong
Use significantly fewer processes for Windows nightlies under pytest.
#32984 opened Oct 29, 2025 by copybara-service bot Loading…
@belitskiy
[doc] specify help string by keyword in configs. documentation kokoro:force-run pull ready Ready for copybara import and testing
#32982 opened Oct 29, 2025 by jakevdp Loading…
@jakevdp
10
Add dynamic grid with smem output test.
#32981 opened Oct 29, 2025 by copybara-service bot Loading…
Use reuse semantics instead of copy semantics when calling DisassembleIntoSingleDeviceArrays in the IFRT proxy, since disassemble always aliases.
#32977 opened Oct 29, 2025 by copybara-service bot Loading…
Rename jax._src.dtypes.dtype to jax._src.dtypes.user_dtype_like_to_dtype.
#32976 opened Oct 29, 2025 by copybara-service bot Loading…
@hawkinsp
Add support for HiType in Pallas Fuser API
#32975 opened Oct 29, 2025 by copybara-service bot Loading…
[pallas:sc] Lower ref indexing to memref.subview instead of tpu.memref_slice and tpu.memref_squeeze
#32973 opened Oct 29, 2025 by copybara-service bot Loading…
@superbobry
#pallas Avoid scaling the pallas_call CostEstimate by the batching dimension size if it is a dynamic shape.
#32971 opened Oct 29, 2025 by copybara-service bot Loading…
Do not error during type promotion when promoting explicit x64 values in a non-x64 mode.
#32967 opened Oct 29, 2025 by copybara-service bot Loading…
@hawkinsp
[pallas:sc] Removed unnecessary libtpu version guards in tests
#32959 opened Oct 29, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Remove WGMMA tests covering large tiles.
#32956 opened Oct 29, 2025 by copybara-service bot Loading…
@bchetioui
Allow type promotion of small integer types.
#32954 opened Oct 29, 2025 by copybara-service bot Loading…
@chr1sj0nes
[Mosaic GPU][NFC] Update the documentation of wgmma.
#32952 opened Oct 29, 2025 by copybara-service bot Loading…
@bchetioui
[Pallas TPU] Automatically plumb name argument to pallas_call into name_stack
#32951 opened Oct 29, 2025 by copybara-service bot Loading…
@sharadmv
[NFC] Add name argument to pallas_call primitive.
#32950 opened Oct 28, 2025 by copybara-service bot Loading…
@sharadmv
[xla:ffi] Migrate to xla::ffi::MakeTypeInfo() API
#32934 opened Oct 28, 2025 by copybara-service bot Loading…
@ezhulenev
[hijax] rename NewstyleHiPrimitive, expose in jax.experimental
#32933 opened Oct 28, 2025 by copybara-service bot Loading…
@mattjj
[CI] Test TSAN job tweaks.
#32932 opened Oct 28, 2025 by copybara-service bot Loading…
@belitskiy
Use a faster pathway for i4 -> f8 canonicalization on TPU 7x: i4 -> i8 -> f8 instead of i4 -> i32 -> f32 -> f8.
#32931 opened Oct 28, 2025 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.