ROCm/jax

Pull Requests Commits

Build with ROCm 7.0 on new runner

charleshofer committed 251 days ago

5b3c49f1

[0.5.0-UT] Fix numerical instability in testEighTinyNorm for ROCm GPU… (#412)

phambinhfin committed 276 days ago

Verified 7fb79ee2

Add workflow for updating XLA commit hash (#330)

charleshofer committed 280 days ago

Verified 6f52ce5a

Skip CSR matmat and matvec float tests on ROCm <6.4 (NaN issue with beta==0) (#384)

psanal35 committed 283 days ago

Verified 5dbfa9d1

Skip ROCm 6.2.4 build for now

charleshofer committed 284 days ago

7c833647

Fix formatting for CI

charleshofer committed 284 days ago

a81c4570

[Pallas TPU] Add better support for BoundedSlice and other BlockDim types to emit_pipeline and fusions

sharadmv committed 284 days ago

ff8341d6

Round-robin pytest tests across GPUs.

hawkinsp committed 284 days ago

f8b77865

Use 32-bit launch ids in `PyLoadedExecutable`

junwhanahn committed 284 days ago

921d1a0d

[JAX] Modify `BatchedCopyToDeviceWithSharding` to use dst shardigns when devices and memory kinds are the same.

ICGog committed 284 days ago

9e1c5b15

chore: switch to pypi hub instead of underlying repos

rickeylev committed 284 days ago

5326d992

Deactivate automatic cluster detection in tests.

mwhittaker committed 284 days ago

24bd9e35

Set PYTHONWARNINGS=error for jax_multiplatform_test.

hawkinsp committed 284 days ago

e7fd33e3

Relax test tolerance for PureCallbackTest.test_can_take_grad_of_pure_callback_with_custom_jvp.

hawkinsp committed 284 days ago

6d53a573

Show Literal's aval when pretty printing.

dfm committed 284 days ago

cbe2b203

Update autodidax.ipynb and autodidax.md

hawkinsp committed 284 days ago

aa7bdd67

PR #28391: Add version guard in autodidax test.

hawkinsp committed 284 days ago

88d6a8fc

Fix test failure in autodidax due to compile() API change.

hawkinsp committed 284 days ago

af336765

SPMD with jax.jit, namedsharding and partition spec

Melody-coder923 committed 284 days ago

ad051116

[JAX] Remove jaxlib/xla directory.

hawkinsp committed 284 days ago

0f2a89dd

Support batch_axis being int rather than Sequence[int] in initializers.

Google-ML-Automation committed 284 days ago

af6d3d70

Fix deprecation warnings in cudnn scaled matmul.

dfm committed 284 days ago

4b21c729

fix: use backend to call xb.process_count in _raise_warnings_or_errors_for_jit_of_pmap

armandpicard committed 284 days ago

7472ea3f

Remove deprecated GPU linalg kernels after compatibility period.

dfm committed 284 days ago

5f15f770

Remove unused parameters in emit_python_callback.

dfm committed 284 days ago

641788a3

Re-land "Inline literals while tracing instead of in a separate pass".

dfm committed 284 days ago

f0a479e1

Update XLA dependency to use revision

Google-ML-Automation committed 284 days ago

65f0e17b

Allow nightly workflows to halt for connection

nitins17 committed 284 days ago

80e89e86

Add note to changelog about nightly package switchover

nitins17 committed 284 days ago

1d76b6d0

[Mosaic GPU] Add suport for packed TMEM

apaszke committed 284 days ago

fdc4209f

Older