Pull requests: jax-ml/jax

New pull request New

562 Open 22,245 Closed

Pull requests list

Use less memory during MLIR compilation.

#32732 opened Oct 18, 2025 by copybara-service bot

Move xla literal out of xla_compiler

#32731 opened Oct 18, 2025 by copybara-service bot

Add logger to libtpu sdk.

#32726 opened Oct 17, 2025 by copybara-service bot

[jax] Migrate JAX to new XLA FFI type registration

#32722 opened Oct 17, 2025 by copybara-service bot

@ezhulenev

[Mosaic] Canonicalize integer tpu::UnpackSubelementsOp with sign_extended=false if it's used by pack that reduces the bitwidth.

#32721 opened Oct 17, 2025 by copybara-service bot

[Pallas:MGPU] PyTorch collectives export WIP

#32719 opened Oct 17, 2025 by copybara-service bot

@apaszke

Fix typo in pjit.py

#32718 opened Oct 17, 2025 by rniczh

[Pallas:MGPU] Substantially refactor the PyTorch conversion logic

#32717 opened Oct 17, 2025 by copybara-service bot

@apaszke

Add a setter for sharding for backward compatibility

#32716 opened Oct 17, 2025 by copybara-service bot

[ROCm] lax.square: use chlo.square kokoro:force-run pull ready

#32715 opened Oct 17, 2025 by nurmukhametov

[pallas:sc] plsc.kernel now supports dict in scratch_shapes=

#32713 opened Oct 17, 2025 by copybara-service bot

@superbobry

Remove unused bias checking code for cuDNN dot_product_attention.

#32711 opened Oct 17, 2025 by copybara-service bot

@chr1sj0nes

[pallas:sc] Run SparseCore tests on presubmit on GitHub

#32696 opened Oct 16, 2025 by copybara-service bot

@superbobry

Optimize multi_weakref_lru_cache.

#32682 opened Oct 16, 2025 by copybara-service bot

@hawkinsp

[NFC][Mosaic:TPU] Expose hash_value of xla::Tile in header

#32681 opened Oct 16, 2025 by copybara-service bot

Add mean argument to jnp.std and jnp.var families

#32668 opened Oct 16, 2025 by arvindsankar

@jakevdp

Remove Traced.in_avals property

#32665 opened Oct 16, 2025 by copybara-service bot

@yashk2810

[Mosaic TPU] Support stochastic rounding on older generations.

#32664 opened Oct 16, 2025 by copybara-service bot

remove literal from xla_compiler

#32662 opened Oct 16, 2025 by copybara-service bot

Reverts 1bd3402c63f3c5c0614af1d732389230d18ec237

#32650 opened Oct 15, 2025 by copybara-service bot

@hawkinsp

[mutable-arrays] enable vjp3 by default

#32635 opened Oct 15, 2025 by copybara-service bot

@mattjj

Use nvml impl lib based wrapper

#32633 opened Oct 15, 2025 by copybara-service bot

@ybaturina

Allow using abstract_mesh in context for trace and lower

#32628 opened Oct 14, 2025 by copybara-service bot

@yashk2810

[xla:ffi] Add TypeRegistry::TypeInfo to be able to register functions to manipulate user-defined types

#32626 opened Oct 14, 2025 by copybara-service bot

@ezhulenev

Exposes jax.Sharding.Index to enable downstream type hints

#32623 opened Oct 14, 2025 by copybara-service bot

ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.

Navigation Menu

Search code, repositories, users, issues, pull requests...

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Pull requests: jax-ml/jax

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Pull requests list