ROCm/jax

Pull Requests Commits

Skip ROCm 6.2.4 build for now

charleshofer committed 261 days ago

7c833647

Fix formatting for CI

charleshofer committed 261 days ago

a81c4570

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

sharadmv committed 261 days ago

ff8341d6

Round-robin pytest tests across GPUs.

hawkinsp committed 261 days ago

f8b77865

Use 32-bit launch ids in `PyLoadedExecutable`

junwhanahn committed 261 days ago

921d1a0d

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

ICGog committed 261 days ago

9e1c5b15

chore: switch to pypi hub instead of underlying repos

rickeylev committed 261 days ago

5326d992

Deactivate automatic cluster detection in tests.

mwhittaker committed 261 days ago

24bd9e35

Set PYTHONWARNINGS=error for jax_multiplatform_test.

hawkinsp committed 261 days ago

e7fd33e3

Relax test tolerance for PureCallbackTest.test_can_take_grad_of_pure_callback_with_custom_jvp.

hawkinsp committed 261 days ago

6d53a573

Show Literal's aval when pretty printing.

dfm committed 261 days ago

cbe2b203

Update autodidax.ipynb and autodidax.md

hawkinsp committed 261 days ago

aa7bdd67

PR #28391: Add version guard in autodidax test.

hawkinsp committed 261 days ago

88d6a8fc

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

hawkinsp committed 261 days ago

af336765

SPMD with jax.jit, namedsharding and partition spec

Melody-coder923 committed 261 days ago

ad051116

[JAX] Remove jaxlib/xla directory.

hawkinsp committed 261 days ago

0f2a89dd

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

Google-ML-Automation committed 261 days ago

af6d3d70

Fix deprecation warnings in cudnn scaled matmul.

dfm committed 261 days ago

4b21c729

fix: use backend to call xb.process_count in _raise_warnings_or_errors_for_jit_of_pmap

armandpicard committed 261 days ago

7472ea3f

Remove deprecated GPU linalg kernels after compatibility period.

dfm committed 261 days ago

5f15f770

Remove unused parameters in emit_python_callback.

dfm committed 261 days ago

641788a3

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

dfm committed 261 days ago

f0a479e1

Update XLA dependency to use revision

Google-ML-Automation committed 261 days ago

65f0e17b

Allow nightly workflows to halt for connection

nitins17 committed 261 days ago

80e89e86

Add note to changelog about nightly package switchover

nitins17 committed 261 days ago

1d76b6d0

[Mosaic GPU] Add suport for packed TMEM

apaszke committed 261 days ago

fdc4209f

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

dimitar-asenov committed 261 days ago

e949ab4c

Add GetAllocatorStats() method for device.

changhuilin committed 261 days ago

ba029782

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

yashk2810 committed 261 days ago

727fe84c

Update JAX nightly index usage

MichaelHudgins committed 261 days ago

a1fe4ca5

Older