ROCm/jax

Pull Requests Commits

Build with ROCm 7.0 on new runner

charleshofer committed 1 year ago

5b3c49f1

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

phambinhfin committed 1 year ago

Verified 7fb79ee2

Add workflow for updating XLA commit hash (#330)

charleshofer committed 1 year ago

Verified 6f52ce5a

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

psanal35 committed 1 year ago

Verified 5dbfa9d1

Skip ROCm 6.2.4 build for now

charleshofer committed 1 year ago

7c833647

Fix formatting for CI

charleshofer committed 1 year ago

a81c4570

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

sharadmv committed 1 year ago

ff8341d6

Round-robin pytest tests across GPUs.

hawkinsp committed 1 year ago

f8b77865

Use 32-bit launch ids in `PyLoadedExecutable`

junwhanahn committed 1 year ago

921d1a0d

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

ICGog committed 1 year ago

9e1c5b15

chore: switch to pypi hub instead of underlying repos

rickeylev committed 1 year ago

5326d992

Deactivate automatic cluster detection in tests.

mwhittaker committed 1 year ago

24bd9e35

Set PYTHONWARNINGS=error for jax_multiplatform_test.

hawkinsp committed 1 year ago

e7fd33e3

Relax test tolerance for PureCallbackTest.test_can_take_grad_of_pure_callback_with_custom_jvp.

hawkinsp committed 1 year ago

6d53a573

Show Literal's aval when pretty printing.

dfm committed 1 year ago

cbe2b203

Update autodidax.ipynb and autodidax.md

hawkinsp committed 1 year ago

aa7bdd67

PR #28391: Add version guard in autodidax test.

hawkinsp committed 1 year ago

88d6a8fc

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

hawkinsp committed 1 year ago

af336765

SPMD with jax.jit, namedsharding and partition spec

Melody-coder923 committed 1 year ago

ad051116

[JAX] Remove jaxlib/xla directory.

hawkinsp committed 1 year ago

0f2a89dd

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

Google-ML-Automation committed 1 year ago

af6d3d70

Fix deprecation warnings in cudnn scaled matmul.

dfm committed 1 year ago

4b21c729

fix: use backend to call xb.process_count in _raise_warnings_or_errors_for_jit_of_pmap

armandpicard committed 1 year ago

7472ea3f

Remove deprecated GPU linalg kernels after compatibility period.

dfm committed 1 year ago

5f15f770

Remove unused parameters in emit_python_callback.

dfm committed 1 year ago

641788a3

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

dfm committed 1 year ago

f0a479e1

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

65f0e17b

Allow nightly workflows to halt for connection

nitins17 committed 1 year ago

80e89e86

Add note to changelog about nightly package switchover

nitins17 committed 1 year ago

1d76b6d0

[Mosaic GPU] Add suport for packed TMEM

apaszke committed 1 year ago

fdc4209f

Older