ROCm/jax

Pull Requests Commits

Use new pool name for RBE

charleshofer committed 2 days ago

a7f81539

Test Bazel with DPX

pelumi1163 committed 2 days ago

aa8b893b

Test Bazel with DPX

pelumi1163 committed 2 days ago

4e1f76fc

Use the new Engflow instance that lives on DOKS (#880)

charleshofer committed 3 days ago

Verified 75b44d43

[ROCm] Run export_harnesses_multi_platform_test_gpu downstream

magaonka-amd committed 7 days ago

e09d7b7b

Merge pull request #40956 from mattjj:push-nyqnoqswnsow

Google-ML-Automation committed 7 days ago

3d4a58a1

[remat3] add custom_vjp.defremat_with_logs

mattjj committed 7 days ago

cd2da890

[remat3] add a remat option to custom_gradient

mattjj committed 7 days ago

8ee4e8eb

[remat3] replace custom_remat with custom_vjp.defremat

mattjj committed 7 days ago

8131b0bc

Split SparseCore Pallas debug check tests into separate targets

bhatuzdaname committed 7 days ago

1b6aae77

JAX changes to support unreduced + partial manual shard_map.

yashk2810 committed 7 days ago

37cd914c

Add unit test verifying ordered jax.debug.callback execution.

hawkinsp committed 7 days ago

4995a7b9

Merge pull request #40876 from jakevdp:cross

Google-ML-Automation committed 8 days ago

4002243d

[pallas:mosaic] Removed `body_prologue=` from the callable returned by `emit_pipeline`

superbobry committed 8 days ago

89060549

[pallas:mosaic] Removed unused argument/return value from `fetch_with_lookahead`

superbobry committed 8 days ago

8e276cbc

[Pallas:TPU] Shortcut lowering of bitcast when src/dst types are equal

apaszke committed 8 days ago

05f18d8b

[pallas:mosaic] A few cosmetic improvements in the pipelining code

superbobry committed 8 days ago

e2fad9a5

[Mosaic GPU] Add layout cast from WGMMA to WGMMA_UPCAST_2X

apaszke committed 8 days ago

4e8a9794

Fix a data race by removing redundant flags.writeable assignment in Array._value.

hawkinsp committed 8 days ago

4e1c02d6

Fix Windows frexp exponent normalization, TPU 7x chip count, and guard recent SparseCore tests.

hawkinsp committed 8 days ago

ef487c41

Removed unused `searchsorted_via_expand`

superbobry committed 8 days ago

c3cca2d0

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

Google-ML-Automation committed 8 days ago

9f01556f

[Mosaic TPU] Separate debug info from the Mosaic kernel `body` into a new `debug_locations` in `custom_call_config` with a new configuration flag `jax_mosaic_unzip_debug_locations`.

yueshengys committed 8 days ago

b1ad9799

Merge pull request #40961 from AshishKumar4:fix-eager-shardmap-check-vma

Google-ML-Automation committed 8 days ago

a3b3b35f

Fix eager partial-manual shard_map with check_vma=False

AshishKumar4 committed 8 days ago

96345404

[pallas] Preserve out_avals memory spaces in `_mpmd_map_discharge_rule`

rdyro committed 8 days ago

024ebf25

Preserve compute_on's SparseCore call by attaching empty backend_config.

Google-ML-Automation committed 8 days ago

9ad1e84d

[Pallas:SC] Assert not compress for unmasked stores

naummo committed 8 days ago

64425ebd

Merge pull request #40933 from mattjj:push-tkywltxmryum

Google-ML-Automation committed 8 days ago

d29c3600

[Pallas:SC] Migrate to expand load ops

naummo committed 8 days ago

cebbb0dd

Older