ROCm/jax

Pull Requests Commits

Remove old scripts and break build for ROCm (#533)

charleshofer committed 92 days ago

47779c38

[0.6.0-UT] Fix ann_test.py::...:test_pmap (Add padding) (#522)

gulsumgudukbay committed 142 days ago

Verified aafa3122

Fix readme and build.py

Ruturaj4 committed 163 days ago

701c3071

docs: add notice to README indicating repo is for internal test purposes only

Ruturaj4 committed 163 days ago

dfc08317

Adjusting block sizes in BWD-FusedAttention tests to fit in GPU shared memory

zahiqbal committed 176 days ago

32c58ad7

Don't build Python 3.10

charleshofer committed 188 days ago

273fa845

Pass shardy option through jax config.

ZixuanJiang committed 189 days ago

b75e2227

add cudnn sdpa mla support

Cjkkkk committed 189 days ago

60d85826

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 189 days ago

00e51508

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 189 days ago

24b0ae00

[JAX] Skip failing tpu tests until June 30th.

subhankarshah committed 189 days ago

57858414

Update jax requirements lock files after 0.6.2 release

kanglant committed 189 days ago

f9bc8652

[doc] fix build error

jakevdp committed 189 days ago

23d99fcb

Fix bugs in the double_buffered_pipeline example

dubstack committed 189 days ago

8e2c2abb

Run pyupgrade --py311-plus.

hawkinsp committed 189 days ago

b2f117e9

Move jax._src.lax to its own BUILD rule

Jake VanderPlas committed 189 days ago

94c1e891

Fix rare error with Literal in DynamicJaxprTracer.full_lower.

jburnim committed 189 days ago

b20b32c4

Skip pytest for Python 3.14 during the JAX release process

kanglant committed 189 days ago

296dd2a7

Remove PositionalSharding from JAX now that 0.6.2 is out and next release is 0.7.0.

yashk2810 committed 189 days ago

2420937e

Remove some dangling references from the docs.

hawkinsp committed 189 days ago

59a26100

[Mosaic GPU] Fix minor error in matmul test.

justinjfu committed 189 days ago

cfdadde2

[Mosaic TPU] Make the backward-compatibility libtpu condition stricter

Google-ML-Automation committed 189 days ago

1a048aa0

Finalize a number of deprecations for JAX v0.7.0

jakevdp committed 189 days ago

8ca8453b

Add a cache around abstract_eval rules.

hawkinsp committed 189 days ago

56cfa598

Remove `_allow_deprecated_jit_signature` now that 0.6.2 is out and next release is 0.7 in July

yashk2810 committed 189 days ago

2746b769

Drop Python 3.10 support.

hawkinsp committed 189 days ago

690bf78a

[Mosaic GPU] Implement canonicalization for `TiledLayout`s.

bchetioui committed 189 days ago

b0a7e8af

[Mosaic GPU] Add support for s8 matmuls on Blackwell

apaszke committed 189 days ago

93d8a323

[Mosaic GPU] Rework the CUDA_ROOT detection once again

apaszke committed 189 days ago

51fea783

[mosaic] Added a `k` prefix to `TPU_MemorySpace` members

superbobry committed 189 days ago

6348ed9f

Older