8000

Commits

Commits on Aug 24, 2026

Skip a Pallas test that is failing on Windows.

hawkinsp
authored and
Google-ML-Automation
committed
8a197ac
[Mosaic:GPU] Detect multi-host mesh in single-controller JAX setup

PatriosTheGreat
authored and
Google-ML-Automation
committed
fbb665d
Upgrade cuDNN version to 9.10.

derdrdirk
authored and
Google-ML-Automation
committed
6f37a70
Add tests for input dilation (LHS dilation / transposed convolution), kernel dilation (RHS dilation / atrous convolution), negative padding (cropping), and strided slicing in Mosaic TPU convolutions.

blakehechtman
authored and
Google-ML-Automation
committed
fc584cd
Merge pull request #40144 from hbahk:fft-gpu-irfft-single-c2r
Google-ML-Automation
committed
3d7c8ca
Merge pull request #40146 from hawkinsp:reqs
Google-ML-Automation
committed
f611dbf
[Mosaic GPU] Saturate OneOf constraints when we fail to infer.

golechwierowicz
authored and
Google-ML-Automation
committed
ce5bab5
Optimize `planar_snake` calculation.

chr1sj0nes
authored and
Google-ML-Automation
committed
6d15f85

Commits on Aug 23, 2026

Fix unsupported gather test not to use a (newly supported) gather.

hawkinsp
authored and
Google-ML-Automation
committed
46eef7d
Regenerate lock files, notably for SciPy 1.18.1.
hawkinsp
committed
dd7f418
Remove dynamic discovery of unknown rocm wheels.

hawkinsp
authored and
Google-ML-Automation
committed
d56fa6a
Update XLA dependency to use revision http://github.com/openxla/xla/commit/e5d008bb03d9b133b0881daf3d108adea8e6625b
Google-ML-Automation
committed
eecec8d
Lower multi-dimensional IRFFT on GPU to a single C2R transform again.
hbahk
committed
2fe0e2b

Commits on Aug 22, 2026

[Mosaic:GPU] Rename multimem parameters mosaic attribute to symmetric memory parameters.

PatriosTheGreat
authored and
Google-ML-Automation
committed
5757465
Update XLA dependency to use revision http://github.com/openxla/xla/commit/ddde0c9ee5a4aa767ff35af8e05473117c7028e4
Google-ML-Automation
committed
6e9c509
Add batching rule for Pallas with_memory_space_constraint primitive.

vanbasten23
authored and
Google-ML-Automation
committed
5e55af7
[Mosaic GPU] Allow creating `MMALayouts` based on `jnp.dtype` as well.

bchetioui
authored and
Google-ML-Automation
committed
3dd8489

Commits on Aug 21, 2026

[Mosaic TPU] Add lowering rule for `lax.top_k` in Pallas Mosaic. For efficiency, we only support `is_stable=False`, i.e., the tie-breaking behavior is non-deterministic.

yueshengys
authored and
Google-ML-Automation
committed
ef14837
Remove temporary UNIMPLEMENTED exception skip guard in `aot_test.py` now that GPU topology deserialization is supported in OpenXLA C-API plugins.

brianseung-google
authored and
Google-ML-Automation
committed
fd8af20
Add machinery for LIT so we can have backend dependent layout rules (pallas TPU, pallas MGPU, JAX HLO, etc).

yashk2810
authored and
Google-ML-Automation
committed
f346f75
Tag a flaky test noasan.

hawkinsp
authored and
Google-ML-Automation
committed
26e24fd
[fuser] Support static argnums and argnames.

rdyro
authored and
Google-ML-Automation
committed
4d3cdf0
[No functional change] Delete VJPHiPrimitive and replace with HiPrim. This is a simple replace.

yashk2810
authored and
Google-ML-Automation
committed
00e27d7
Attempt at fixing a rare flake in pytest CI.

hawkinsp
authored and
Google-ML-Automation
committed
b5dbb67
Use PORTSERVER_ADDRESS to communicate with CI jobs, not JAX_PORTSERVER_ADDRESS.

hawkinsp
authored and
Google-ML-Automation
committed
fcf7786
[hijax][custom_vjp3] skip sharding canonicalization and preserve effects in call_hi_primitive and eval_jaxpr_p

mattjj
authored and
Google-ML-Automation
committed
d3b67fc
Roll back PR #40111: break tests.
Google-ML-Automation
committed
cc882a4
Fix some tests flaking in CI.

hawkinsp
authored and
Google-ML-Automation
committed
4ef3f19
Update XLA dependency to use revision http://github.com/openxla/xla/commit/3abee12b628f5ba4fcb8d8f849e3dda68b801ecf
Google-ML-Automation
committed
3a32b56
Rename `VJPHiPrimitive` to `HiPrim` but keep `VJPHiPrimitive` as an alias

yashk2810
authored and
Google-ML-Automation
committed
4fe874b
Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive

yashk2810
authored and
Google-ML-Automation
committed
c90c245
[pltpu] Support allocations in the emit_pipeline primitive.

rdyro
authored and
Google-ML-Automation
committed
6ff936b
Disallow closing over Refs with manual axis types in mpmd_map

sharadmv
authored and
Google-ML-Automation
committed
9548051
Fix incompatibilities with nanobind v2.14.0

hawkinsp
authored and
Google-ML-Automation
committed
2cc327d

Commits on Aug 20, 2026

Merge pull request #40114 from mattjj:push-prnmvzqoxlzr
Google-ML-Automation
committed
02f1f0f

0