ROCm/jax

Pull Requests Commits

Fix print statement syntax for ROCm error message (#574)

gulsumgudukbay committed 190 days ago

Verified 9f5ee019

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

charleshofer committed 288 days ago

Verified 2da601b9

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

gulsumgudukbay committed 318 days ago

Verified aafa3122

Fix readme and build.py

Ruturaj4 committed 339 days ago

701c3071

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

Ruturaj4 committed 340 days ago

dfc08317

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

zahiqbal committed 353 days ago

32c58ad7

Don't build Python 3.10

charleshofer committed 1 year ago

273fa845

Pass shardy option through jax config.

ZixuanJiang committed 1 year ago

b75e2227

add cudnn sdpa mla support

Cjkkkk committed 1 year ago

60d85826

[Pallas][Mosaic GPU] Add GPU pipelining docs

justinjfu committed 1 year ago

00e51508

Add wrap_negative_indices paramter to jnp.ndarray.at[]

jakevdp committed 1 year ago

24b0ae00

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

subhankarshah committed 1 year ago

57858414

Update jax requirements lock files after 0.6.2 release

kanglant committed 1 year ago

f9bc8652

[doc] fix build error

jakevdp committed 1 year ago

23d99fcb

Fix bugs in the double_buffered_pipeline example

dubstack committed 1 year ago

8e2c2abb

Run pyupgrade --py311-plus.

hawkinsp committed 1 year ago

b2f117e9

Move jax._src.lax to its own BUILD rule

Jake VanderPlas committed 1 year ago

94c1e891

Fix rare error with Literal in DynamicJaxprTracer.full_lower.

jburnim committed 1 year ago

b20b32c4

Skip pytest for Python 3.14 during the JAX release process

kanglant committed 1 year ago

296dd2a7

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

yashk2810 committed 1 year ago

2420937e

Remove some dangling references from the docs.

hawkinsp committed 1 year ago

59a26100

[Mosaic GPU] Fix minor error in matmul test.

justinjfu committed 1 year ago

cfdadde2

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

Google-ML-Automation committed 1 year ago

1a048aa0

Finalize a number of deprecations for JAX v0.7.0

jakevdp committed 1 year ago

8ca8453b

Add a cache around abstract_eval rules.

hawkinsp committed 1 year ago

56cfa598

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

yashk2810 committed 1 year ago

2746b769

Drop Python 3.10 support.

hawkinsp committed 1 year ago

690bf78a

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

bchetioui committed 1 year ago

b0a7e8af

[Mosaic GPU] Add support for s8 matmuls on Blackwell

apaszke committed 1 year ago

93d8a323

[Mosaic GPU] Rework the CUDA_ROOT detection once again

apaszke committed 1 year ago

51fea783

Older