ROCm/jax

Pull Requests Commits

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

charleshofer committed 156 days ago

47779c38

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

gulsumgudukbay committed 206 days ago

Verified aafa3122

Fix readme and build.py

Ruturaj4 committed 227 days ago

701c3071

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

Ruturaj4 committed 228 days ago

dfc08317

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

zahiqbal committed 241 days ago

32c58ad7

Don't build Python 3.10

charleshofer committed 253 days ago

273fa845

Pass shardy option through jax config.

ZixuanJiang committed 253 days ago

b75e2227

add cudnn sdpa mla support

Cjkkkk committed 253 days ago

60d85826

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 253 days ago

00e51508

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 253 days ago

24b0ae00

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

subhankarshah committed 253 days ago

57858414

Update jax requirements lock files after 0.6.2 release

kanglant committed 253 days ago

f9bc8652

[doc] fix build error

jakevdp committed 253 days ago

23d99fcb

Fix bugs in the double_buffered_pipeline example

dubstack committed 253 days ago

8e2c2abb

Run pyupgrade --py311-plus.

hawkinsp committed 253 days ago

b2f117e9

Move jax._src.lax to its own BUILD rule

Jake VanderPlas committed 253 days ago

94c1e891

Fix rare error with Literal in DynamicJaxprTracer.full_lower.

jburnim committed 253 days ago

b20b32c4

Skip pytest for Python 3.14 during the JAX release process

kanglant committed 253 days ago

296dd2a7

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

yashk2810 committed 253 days ago

2420937e

Remove some dangling references from the docs.

hawkinsp committed 253 days ago

59a26100

[Mosaic GPU] Fix minor error in matmul test.

justinjfu committed 253 days ago

cfdadde2

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

Google-ML-Automation committed 253 days ago

1a048aa0

Finalize a number of deprecations for JAX v0.7.0

jakevdp committed 253 days ago

8ca8453b

Add a cache around abstract_eval rules.

hawkinsp committed 253 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 253 days ago

2746b769

Drop Python 3.10 support.

hawkinsp committed 253 days ago

690bf78a

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

bchetioui committed 253 days ago

b0a7e8af

[Mosaic GPU] Add support for s8 matmuls on Blackwell

apaszke committed 253 days ago

93d8a323

[Mosaic GPU] Rework the CUDA_ROOT detection once again

apaszke committed 253 days ago

51fea783

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

superbobry committed 253 days ago

6348ed9f

Older