ROCm/jax

Pull Requests Commits

Split the multi-accelerator ROCm tests into their own script

gulsumgudukbay committed 12 days ago

6b6b8b56

[ROCm] Simplify pytest parallelism calculation for ROCm CI

magaonka-amd committed 12 days ago

ce673d88

Add emulation for int4 cmpi

Sam Ginzburg committed 12 days ago

798b72fa

Merge pull request #39915 from mattjj:push-ustwymvzzpon

Google-ML-Automation committed 12 days ago

da2d20de

Make dce_sink a VJPHiPrimitive so that it works with sharding

yashk2810 committed 12 days ago

4b315c97

[custom_vjp3] fix export + gspmdsharding + custom_vjp3

mattjj committed 12 days ago

2f55e320

Merge pull request #39606 from Intel-tensorflow:oneapi_callback_support

Google-ML-Automation committed 12 days ago

5a55d303

Merge pull request #39912 from mattjj:push-vrrrxxwtyyor

Google-ML-Automation committed 13 days ago

cc77b3bb

[custom_vjp3] misc fixes, in-tree tests now green

mattjj committed 13 days ago

27b42375

Merge pull request #39765 from ROCm:drop-rocm-sdk-path-block-upstream

Google-ML-Automation committed 13 days ago

d46fcd96

Fix ASan memory leak in PJRT custom call handler registration

akuegel committed 13 days ago

90eac75f

[pallas:mgpu] Use `mpmd_map` instead of `core_map` in `plgpu.warp_map`

superbobry committed 13 days ago

79616548

Merge pull request #39909 from mattjj:push-wpyxnttptsxm

Google-ML-Automation committed 13 days ago

593a0203

[remat3] fix vmap axes=None bug

mattjj committed 13 days ago

a6e59252

Merge pull request #39340 from wangfakang:opt_topo_aware

Google-ML-Automation committed 13 days ago

99783a3a

Wrap negative indices by default in take_along_axis and remove legacy behavior that used promise_in_bounds.

18praveenb committed 13 days ago

bbf826af

add python and buffer callback support for oneapi

kranipa committed 13 days ago

1f3308d0

[pallas] Migrated a few more test usages away from `pl.core_map`

superbobry committed 13 days ago

003fbe7f

Update XLA dependency to use revision

Google-ML-Automation committed 13 days ago

27b2af44

Add `gpu_a100` to `mosaic_gpu_test`'s `enable_configs`.

bchetioui committed 13 days ago

5dc3e31e

[Pallas/Mosaic GPU] Fix integer dtype mismatches in `emit_pipeline`.

bchetioui committed 13 days ago

8730daa4

[Pallas/Mosaic GPU] Remove `notap` tag from `mgpu_examples_test`.

bchetioui committed 13 days ago

264d0b4e

[Pallas/Mosaic GPU] Fix usage of `pl.warp_map` in `mgpu_examples_test`

bchetioui committed 13 days ago

a1a49e3e

[Mosaic GPU] Better error message for unsatisfiable constraint system for dynamic indices.

golechwierowicz committed 13 days ago

f4a195e8

Merge pull request #39890 from jax-ml:warn

Google-ML-Automation committed 13 days ago

248a4a00

Set incompatible_default_to_explicit_init_py in .bazelrc

hawkinsp committed 13 days ago

79a900ef

Merge pull request #39882 from jax-ml:locks

Google-ML-Automation committed 13 days ago

e0c3dc62

Bump the XLA commit and remove the temporary rules_python_220.patch.

hawkinsp committed 13 days ago

fe0f1193

Bump rules_python to v2.2.0.

hawkinsp committed 13 days ago

6d280128

Consolidate freethreading and non-freethreading requirements lock files.

hawkinsp committed 13 days ago

4203b615

Older