ROCm/jax

Pull Requests Commits

[ROCm] Switch CI from ROCm 7.14 to ROCm 10.0.0rc4

magaonka-amd committed 14 hours ago

1c10c683

[ROCm] Enable buffer_callback_test_gpu in ROCm CI

magaonka-amd committed 14 hours ago

7af04405

[ROCm] Enable the 35 device-agnostic bazel tests in ROCm CI

magaonka-amd committed 14 hours ago

d2bac59c

[ROCm] continuous: build jax/jaxlib on fork runners

magaonka-amd committed 14 hours ago

c6249638

[ROCm] nightly: build jax/jaxlib on fork runners

mminutoli committed 14 hours ago

368308ee

Enable additional single-GPU bazel tests in ROCm CI (#816)

magaonka-amd committed 14 hours ago

65502bdc

[pallas:mgpu] Allowed specifying `m_warps=` in the MMA_* layouts

superbobry committed 14 hours ago

ead0c19c

[pallas:sc] Fixed effect propagation for closed-over refs in `plsc.parallel_loop`

superbobry committed 14 hours ago

5daac240

Skip a Pallas test that is failing on Windows.

hawkinsp committed 16 hours ago

8a197ac0

[Mosaic:GPU] Detect multi-host mesh in single-controller JAX setup

PatriosTheGreat committed 17 hours ago

fbb665d7

Upgrade cuDNN version to 9.10.

derdrdirk committed 17 hours ago

6f37a700

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 committed 18 hours ago

fc584cdf

Merge pull request #40144 from hbahk:fft-gpu-irfft-single-c2r

Google-ML-Automation committed 18 hours ago

3d7c8ca9

Merge pull request #40146 from hawkinsp:reqs

Google-ML-Automation committed 18 hours ago

f611dbf3

[Mosaic GPU] Saturate OneOf constraints when we fail to infer.

golechwierowicz committed 21 hours ago

ce5bab5b

Optimize `planar_snake` calculation.

chr1sj0nes committed 23 hours ago

6d15f85f

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

hawkinsp committed 1 day ago

46eef7d1

Regenerate lock files, notably for SciPy 1.18.1.

hawkinsp committed 1 day ago

dd7f4185

Remove dynamic discovery of unknown rocm wheels.

hawkinsp committed 1 day ago

d56fa6ae

Update XLA dependency to use revision http://github.com/openxla/xla/commit/e5d008bb03d9b133b0881daf3d108adea8e6625b

Google-ML-Automation committed 1 day ago

eecec8db

Lower multi-dimensional IRFFT on GPU to a single C2R transform again.

hbahk committed 2 days ago

2fe0e2b2

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

PatriosTheGreat committed 2 days ago

5757465a

Update XLA dependency to use revision http://github.com/openxla/xla/commit/ddde0c9ee5a4aa767ff35af8e05473117c7028e4

Google-ML-Automation committed 2 days ago

6e9c5090

Add batching rule for Pallas with_memory_space_constraint primitive.

vanbasten23 committed 3 days ago

5e55af7f

[Mosaic GPU] Allow creating `MMALayouts` based on `jnp.dtype` as well.

bchetioui committed 3 days ago

3dd8489c

[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 committed 3 days ago

ef148373

Remove temporary UNIMPLEMENTED exception skip guard in `aot_test.py` now that GPU topology deserialization is supported in OpenXLA C-API plugins.

brianseung-google committed 3 days ago

fd8af205

Add machinery for LIT so we can have backend dependent layout rules (pallas TPU, pallas MGPU, JAX HLO, etc).

yashk2810 committed 3 days ago

f346f754

Tag a flaky test noasan.

hawkinsp committed 3 days ago

26e24fd8

[fuser] Support static argnums and argnames.

rdyro committed 3 days ago

4d3cdf04

Older