ROCm/jax

Pull Requests Commits

[ROCm] Run ROCm pytest with 4 workers per GPU

magaonka-amd committed 5 days ago

024d6cf8

[ROCm] nightly: run only the ROCm wheel build and 1-GPU pytest

magaonka-amd committed 5 days ago

c1c6ca99

[ROCm] nightly: run only the ROCm wheel build and 4-GPU pytest

magaonka-amd committed 7 days ago

2c9c7fcc

[ROCm] Stop running the continuous wheel tests on a schedule

magaonka-amd committed 7 days ago

8e1d2b5f

[ROCm] Run export_harnesses_multi_platform_test_gpu downstream

magaonka-amd committed 7 days ago

eed75d72

[JAX - Pallas] Skip `test_pltpu_bitcast_minormost` when `libtpu < 0.0.50`.

belitskiy committed 7 days ago

8eb6aff6

[pallas/fuser] Improve custom_fusion validation and error messages based on the TODO comments

IvyZX committed 7 days ago

516d99a5

[Mosaic GPU][NFC] Raise error message in units of bytes instead of bits.

bchetioui committed 7 days ago

c5df14c8

Change `ragged_dot` `group_offset` to a scalar array.

chr1sj0nes committed 7 days ago

99dd9e4c

[Pallas:SC] Add dim to pltpu.bitcast and delete plsc.bitcast_p

naummo committed 7 days ago

7a1535cc

Improve erf_inv accuracy near |x| = 1.

hawkinsp committed 7 days ago

4bbb23df

Improve float32 log1p accuracy on TPU v5p, v6e and 7x.

hawkinsp committed 7 days ago

11fa76a9

[Mosaic GPU] Delete `WaitOp` from the dialect.

bchetioui committed 7 days ago

b0903459

[Mosaic GPU][NFC] Pre-factor out `before_thread_sync` helper.

bchetioui committed 8 days ago

c12ad12f

[mgpu] Add barrier_arrive_and_wait primitive.

chr1sj0nes committed 8 days ago

4154b4a0

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

Google-ML-Automation committed 8 days ago

bb89a781

[pallas] Fix emit_pipeline under jax_enable_x64

IvyZX committed 8 days ago

1e0af75c

[Pallas SC] Support remote transfers between HBM and vector subcore VMEM (rewrite to shared VMEM).

emilyfertig committed 8 days ago

b3b7d2e3

Merge pull request #41302 from P-bibs:lazy-barrier-dealloc-check

Google-ML-Automation committed 8 days ago

8ed533cf

fix instances of "the the" in JAX

Jake VanderPlas committed 8 days ago

1584d6dd

Reimplement `jax.scipy.special.erfcx` and add `accuracy=HIGHEST` for `lax.erfc`, with accurate derivatives.

hawkinsp committed 8 days ago

e9b14a6f

Merge pull request #41276 from Kuldeeep18:fix/norm-isf-precision

Google-ML-Automation committed 8 days ago

de856c17

Merge pull request #40788 from mayuriphad:fix-callback-error-msg

Google-ML-Automation committed 8 days ago

c24f87b6

[Mosaic TPU] Extend support of reshapes of form (*src_major, src_minor) <-> (*dst_major, dst_minor), where max(src_minor, dst_minor) <= 128 and one divides the other (row merging and splitting) to 1D cases.

yueshengys committed 8 days ago

40fc7a89

Merge pull request #41324 from jakevdp:astype-annotation

Google-ML-Automation committed 8 days ago

fc8712e0

Merge pull request #41233 from ROCm:cph-upstream-use-new-rbe-cluster

Google-ML-Automation committed 8 days ago

3d87ca7b

Replace `PyTuple_SET_ITEM` and `PyList_SET_ITEM` with nanobind builders.

hawkinsp committed 8 days ago

df3e275a

[typing] fix jax.Array.astype annotation

jakevdp committed 8 days ago

06667818

[Mosaic GPU] Narrow `collective` type in `LaunchContext.async_copy` and `async_prefetch` to use a `tuple` instead of a `Sequence`.

bchetioui committed 8 days ago

922602f3

[CI] Restrict `google-benchmark` to `python_version<"3.14"` in `build/test-requirements.txt`.

belitskiy committed 8 days ago

d2357e63

Older