ROCm/jax

Pull Requests Commits

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

zahiqbal committed 253 days ago

4757f774

Don't build Python 3.10

charleshofer committed 259 days ago

273fa845

Pass shardy option through jax config.

ZixuanJiang committed 260 days ago

b75e2227

add cudnn sdpa mla support

Cjkkkk committed 260 days ago

60d85826

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 260 days ago

00e51508

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 260 days ago

24b0ae00

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

subhankarshah committed 260 days ago

57858414

Update jax requirements lock files after 0.6.2 release

kanglant committed 260 days ago

f9bc8652

[doc] fix build error

jakevdp committed 260 days ago

23d99fcb

Fix bugs in the double_buffered_pipeline example

dubstack committed 260 days ago

8e2c2abb

Run pyupgrade --py311-plus.

hawkinsp committed 260 days ago

b2f117e9

Move jax._src.lax to its own BUILD rule

Jake VanderPlas committed 260 days ago

94c1e891

Fix rare error with Literal in DynamicJaxprTracer.full_lower.

jburnim committed 260 days ago

b20b32c4

Skip pytest for Python 3.14 during the JAX release process

kanglant committed 260 days ago

296dd2a7

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

yashk2810 committed 260 days ago

2420937e

Remove some dangling references from the docs.

hawkinsp committed 260 days ago

59a26100

[Mosaic GPU] Fix minor error in matmul test.

justinjfu committed 260 days ago

cfdadde2

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

Google-ML-Automation committed 260 days ago

1a048aa0

Finalize a number of deprecations for JAX v0.7.0

jakevdp committed 260 days ago

8ca8453b

Add a cache around abstract_eval rules.

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

2746b769

Drop Python 3.10 support.

hawkinsp committed 260 days ago

690bf78a

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

bchetioui committed 260 days ago

b0a7e8af

[Mosaic GPU] Add support for s8 matmuls on Blackwell

apaszke committed 260 days ago

93d8a323

[Mosaic GPU] Rework the CUDA_ROOT detection once again

apaszke committed 260 days ago

51fea783

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

superbobry committed 260 days ago

6348ed9f

[Mosaic:TPU] Byte-granularity dynamic gathers

tlongeri committed 260 days ago

66795a13

jaxlib_extension_version == 355 after 0.6.2 release. So remove the conditionals.

yashk2810 committed 260 days ago

361cf02d

Add `gather` to `roofline`.

zacmustin committed 260 days ago

4fb6f1bc

Postrelease (0.6.2) changes

yashk2810 committed 260 days ago

98e91122

Older