ROCm/jax

Pull Requests Commits

Skip test_tridiagonal_solve_grad0 on AMD due to rocSparse numerical issue

AratiGanesh committed 53 days ago

9198a10b

Update XLA dependency to use revision http://github.com/openxla/xla/commit/52f095f542408e774ee5a847873b0a497bdeb873

Google-ML-Automation committed 53 days ago

8304fec2

Merge pull request #34332 from jakevdp:searchsorted-hijax

Google-ML-Automation committed 54 days ago

13dcc922

Reduce the number of pytest workers for Linux aarch64 CPU builds.

hawkinsp committed 54 days ago

6463906c

Use a np.printoptions context manager rather than np.set_printoptions in test utility.

hawkinsp committed 54 days ago

5c4dd597

Reverts 3b4bee803de1a83b98cfe1aced47ab58f9bbc0c3

hawkinsp committed 54 days ago

7a68c4b1

jakevdp committed 54 days ago

5d7b5186

validate out_dtype

jakevdp committed 54 days ago

250ced1d

more fixes for sharded computation

jakevdp committed 54 days ago

aad9769e

Use scan rather than fori_loop

jakevdp committed 54 days ago

6ef7e8d8

[hijax] add hijax searchsorted primitive

jakevdp committed 54 days ago

2185cb1c

Update XLA dependency to use revision http://github.com/openxla/xla/commit/ac24cc91d12634f06c21a30f170c6a117ff615f9

Google-ML-Automation committed 54 days ago

4ea536c2

Merge pull request #35696 from davisyoshida:recip-docstring

Google-ML-Automation committed 54 days ago

ea4b2e7e

davisyoshida committed 54 days ago

a767ad38

Fix reciprocal docstring

davisyoshida committed 54 days ago

f5676668

Automated Code Change

Google-ML-Automation committed 55 days ago

2d5a40fe

Make _get_owner skip sharding_constraints so that the xla_metadata is applied to the op instead of the sharding constraints

yashk2810 committed 55 days ago

a364e282

[Take 2] shard map without linear util

dougalm committed 55 days ago

830e342c

Perform dtype canonicalization during Primitive.bind.

hawkinsp committed 55 days ago

fae1e236

Merge pull request #35818 from jakevdp:lax-accuracy

Google-ML-Automation committed 55 days ago

9915f8c2

Merge pull request #35852 from ROCm:jax-rocm-version

Google-ML-Automation committed 55 days ago

964e2ebc

Merge pull request #35866 from jakevdp:take-current-trace

Google-ML-Automation committed 55 days ago

a8da2f48

Increase size of test jet_test_cpu to avoid timeout.

mwhittaker committed 55 days ago

326f8b44

Improve handling of accuracy parameter in jax.lax

jakevdp committed 55 days ago

fa2a8caf

Reverts 0b56550381c1f5648235606a92239c12bd1e876f

Google-ML-Automation committed 55 days ago

8c5def0f

[Pallas] Add buffering mode to pltpu that enables non-immediate revisit

Google-ML-Automation committed 55 days ago

1b268ffc

Merge pull request #35894 from hawkinsp:docports

Google-ML-Automation committed 55 days ago

573b73d7

Attempt 2 at fixing doc build port collisions.

hawkinsp committed 55 days ago

3e463e76

add jax[rocm-local] extra and relax plugin version check for accepting wildcards

gulsumgudukbay committed 55 days ago

2fb2d6cf

[Mosaic GPU] Fix a race condition in the 2CTA SMEM->TMEM copy test

apaszke committed 55 days ago

4acb04de

Older