ROCm/jax

Pull Requests Commits

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

[Mosaic GPU] Add layout inference and lowering for `vector.MultiDimReduce`

dimitar-asenov committed 1 year ago

e949ab4c

Add GetAllocatorStats() method for device.

changhuilin committed 1 year ago

ba029782

Enter into correct mesh context in shard_map in `_partial_eval_jaxpr_custom_rule`.

yashk2810 committed 1 year ago

727fe84c

Update JAX nightly index usage

MichaelHudgins committed 1 year ago

a1fe4ca5

Older