Skip to content

Pull requests: jax-ml/jax

New pull request New
558 Open 21,924 Closed
558 Open 21,924 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

Make convert_element_type_p a standard primitive and remove the usage of utils.standard_abstract_eval.
#32348 opened Oct 4, 2025 by copybara-service bot Loading…
@yashk2810
Update JAX get_tfrt_cpu_client to use GetPjRtCpuClient directly.
#32344 opened Oct 3, 2025 by copybara-service bot Loading…
[XLA:Python] Fix Python profiler hooks under Python free threading.
#32341 opened Oct 3, 2025 by copybara-service bot Loading…
@hawkinsp
jnp.setdiff1d: handle non-1D inputs with assume_unique=True kokoro:force-run pull ready Ready for copybara import and testing
#32340 opened Oct 3, 2025 by jakevdp Loading…
1
@jakevdp
[Mosaic TPU][relayout] Mask bitwidth change with any change of tiling, offset and implicit_dim
#32332 opened Oct 3, 2025 by copybara-service bot Loading…
Support For Roofline Analysis Of Pallas Kernels
#32328 opened Oct 3, 2025 by copybara-service bot Loading…
[Mosaic TPU] Trace all basic blocks by default
#32327 opened Oct 3, 2025 by copybara-service bot Loading…
Support For Roofline Analysis Of Pallas Kernels
#32325 opened Oct 3, 2025 by copybara-service bot Loading…
Update rules_ml_toolchain and set LLVM version to 21.
#32323 opened Oct 3, 2025 by copybara-service bot Loading…
[JAX] Add require_user_context_scope config in PjRt-IFRT and enable it for JAX
#32321 opened Oct 3, 2025 by copybara-service bot Loading…
[Pallas] Add HiType support for pallas_call
#32320 opened Oct 3, 2025 by copybara-service bot Loading…
@sharadmv
Make jax-abstract-dunder-array a hard error.
#32318 opened Oct 2, 2025 by copybara-service bot Loading…
@hawkinsp
[JAX] Add compiled program call_location logging support to JAX
#32316 opened Oct 2, 2025 by copybara-service bot Loading…
Use slice instead of iota for simple cases of vmap over Refs.
#32315 opened Oct 2, 2025 by copybara-service bot Loading…
@justinjfu
Set optimize_scratch_bytes to false in Fusion Autotuner Pass
#32314 opened Oct 2, 2025 by copybara-service bot Loading…
Relax tolerance for test_emit_with_no_output. When setting --xla_gpu_experimental_enable_fusion_autotuner, this test will become flaky when it goes through the BlockLevelEmitter. It looks like the tolerance is too high because the differences seem quite small and the autotuner doesn't mark it as having incorrect results, but let me know if there's more I should check.
#32313 opened Oct 2, 2025 by copybara-service bot Loading…
Autotune fusions through triton & native emitters
#32312 opened Oct 2, 2025 by copybara-service bot Loading…
[Mosaic GPU] Compress the representation of Divides constraints.
#32299 opened Oct 2, 2025 by copybara-service bot Loading…
@bchetioui
[Mosaic TPU] Relax second minor dim shape restrictions for packed second/third minor transpose
#32297 opened Oct 2, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Add limited support for third/second minor transposes of packed types
#32295 opened Oct 2, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Clearly define the semantics of interleaved packing
#32294 opened Oct 2, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Only canonicalize the dtype of transposes if they use the XLU
#32290 opened Oct 2, 2025 by copybara-service bot Loading…
@apaszke
[jaxlib] Generate .pyi files for the native extensions via nanobind stubgen
#32287 opened Oct 2, 2025 by copybara-service bot Loading…
@superbobry
docs(pallas): Fix broken link to pallas documentation section
#32284 opened Oct 2, 2025 by nicholasjng Loading…
Support For Roofline Analysis Of Pallas Kernels
#32281 opened Oct 2, 2025 by copybara-service bot Loading…
ProTip! Mix and match filters to narrow down what you’re looking for.