Skip to content

Pull requests: jax-ml/jax

New pull request New
587 Open 19,702 Closed
587 Open 19,702 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

Add out_sharding to jnp.ones, jnp.zeros and jnp.emtpy
#29814 opened Jun 27, 2025 by copybara-service bot Loading…
@yashk2810
Replace old heartbeat options with simplified heartbeat options.
#29813 opened Jun 27, 2025 by copybara-service bot Loading…
@mwhittaker
Populate the HLO op_type field from JAX StableHLO.
#29811 opened Jun 27, 2025 by copybara-service bot Loading…
@hawkinsp
Replace weakref_lru_cache with a weakref.WeakValueDictionary kokoro:force-run pull ready Ready for copybara import and testing
#29808 opened Jun 27, 2025 by gnecula Loading…
1
5
[Pallas:TPU] Add a missing check for axis being within bounds for pltpu.repeat
#29806 opened Jun 27, 2025 by copybara-service bot Loading…
@apaszke
[pallas] Generalize support for merging dimensions in fuser reshape.
#29800 opened Jun 27, 2025 by copybara-service bot Loading…
@chr1sj0nes
[Mosaic:TPU][Relayouts] Fallback for implicit dim change that goes through 32-bit native tiling
#29799 opened Jun 27, 2025 by copybara-service bot Loading…
[Mosaic:TPU][Relayouts] 2nd minor implicit -> minor implicit
#29795 opened Jun 26, 2025 by copybara-service bot Loading…
Include dispatch primitives in roofline.
#29793 opened Jun 26, 2025 by copybara-service bot Loading…
[JAX] Add a cache around HLO lowering rules.
#29789 opened Jun 26, 2025 by copybara-service bot Loading…
@hawkinsp
[Mosaic] Pass rounding mode to truncf.
#29787 opened Jun 26, 2025 by copybara-service bot Loading…
Remove JAX_IFRT_VERSION_NUMBER check for types in XLA/JAX
#29784 opened Jun 26, 2025 by copybara-service bot Loading…
Do not propagate name_stacks into lower_jaxpr_to_fun.
#29783 opened Jun 26, 2025 by copybara-service bot Loading…
@hawkinsp
[Mosaic GPU] Add an __or__ implementation for EquationSystem.
#29777 opened Jun 26, 2025 by copybara-service bot Loading…
@bchetioui
[Mosaic GPU] Add an initial skeleton for an equation system for layout/transform inference.
#29771 opened Jun 26, 2025 by copybara-service bot Loading…
@bchetioui
Expose DeviceList in xla_client.
#29767 opened Jun 26, 2025 by copybara-service bot Loading…
Add memory_space to state.AbstractRef and use it instead of AbstractMemoryRef in Pallas
#29761 opened Jun 26, 2025 by copybara-service bot Loading…
@sharadmv
[Mosaic:TPU][Relayouts] Materialize replicated offsets
#29760 opened Jun 25, 2025 by copybara-service bot Loading…
[Mosaic] Extend indirect DMA verification to dynamic shapes
#29759 opened Jun 25, 2025 by copybara-service bot Loading…
[Mosaic] Add DMA source and target shape verification
#29755 opened Jun 25, 2025 by copybara-service bot Loading…
[Pallas:MGPU] Make the optimized flag on load/stores a compiler param
#29748 opened Jun 25, 2025 by copybara-service bot Loading…
@apaszke
[jaxpr consts] Propagate closed-over constants to the top-level jit pull ready Ready for copybara import and testing
#29743 opened Jun 25, 2025 by gnecula Loading…
@gnecula
[pallas] Moved pl.broadcast_to into a Triton-specific submodule
#29736 opened Jun 25, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic:TPU][Relayout] (1, 128 * packing) <-> (n * packing, 128) retilings for n > 1
#29732 opened Jun 25, 2025 by copybara-service bot Loading…
Add [query,key_value]_seq_offsets arguments to dot_product_attention
#29731 opened Jun 25, 2025 by chaserileyroberts Loading…
1
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.