ROCm/jax

Pull Requests Commits

The reshape call has a numerical bug where the target reshape size is always higher than array size if num_devices is not a nonnegative power of two. To fix this, this commit adds padding

gulsumgudukbay committed 217 days ago

26d0e840

Fix readme and build.py

Ruturaj4 committed 232 days ago

701c3071

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

Ruturaj4 committed 232 days ago

dfc08317

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

zahiqbal committed 246 days ago

32c58ad7

Don't build Python 3.10

charleshofer committed 258 days ago

273fa845

Pass shardy option through jax config.

ZixuanJiang committed 258 days ago

b75e2227

add cudnn sdpa mla support

Cjkkkk committed 258 days ago

60d85826

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 258 days ago

00e51508

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 258 days ago

24b0ae00

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

subhankarshah committed 258 days ago

57858414

Update jax requirements lock files after 0.6.2 release

kanglant committed 258 days ago

f9bc8652

[doc] fix build error

jakevdp committed 258 days ago

23d99fcb

Fix bugs in the double_buffered_pipeline example

dubstack committed 258 days ago

8e2c2abb

Run pyupgrade --py311-plus.

hawkinsp committed 258 days ago

b2f117e9

Move jax._src.lax to its own BUILD rule

Jake VanderPlas committed 258 days ago

94c1e891

Fix rare error with Literal in DynamicJaxprTracer.full_lower.

jburnim committed 258 days ago

b20b32c4

Skip pytest for Python 3.14 during the JAX release process

kanglant committed 258 days ago

296dd2a7

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

yashk2810 committed 258 days ago

2420937e

Remove some dangling references from the docs.

hawkinsp committed 258 days ago

59a26100

[Mosaic GPU] Fix minor error in matmul test.

justinjfu committed 258 days ago

cfdadde2

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

Google-ML-Automation committed 258 days ago

1a048aa0

Finalize a number of deprecations for JAX v0.7.0

jakevdp committed 258 days ago

8ca8453b

Add a cache around abstract_eval rules.

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

2746b769

Drop Python 3.10 support.

hawkinsp committed 258 days ago

690bf78a

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

bchetioui committed 258 days ago

b0a7e8af

[Mosaic GPU] Add support for s8 matmuls on Blackwell

apaszke committed 258 days ago

93d8a323

[Mosaic GPU] Rework the CUDA_ROOT detection once again

apaszke committed 258 days ago

51fea783

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

superbobry committed 258 days ago

6348ed9f

[Mosaic:TPU] Byte-granularity dynamic gathers

tlongeri committed 258 days ago

66795a13

Older