ROCm/jax

Pull Requests Commits

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

Google-ML-Automation committed 73 days ago

8304fec2

Merge pull request #34332 from jakevdp:searchsorted-hijax

Google-ML-Automation committed 74 days ago

13dcc922

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

hawkinsp committed 74 days ago

6463906c

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

hawkinsp committed 74 days ago

5c4dd597

Reverts 3b4bee803de1a83b98cfe1aced47ab58f9bbc0c3

hawkinsp committed 74 days ago

7a68c4b1

jakevdp committed 74 days ago

5d7b5186

validate out_dtype

jakevdp committed 74 days ago

250ced1d

more fixes for sharded computation

jakevdp committed 74 days ago

aad9769e

Use scan rather than fori_loop

jakevdp committed 74 days ago

6ef7e8d8

[hijax] add hijax searchsorted primitive

jakevdp committed 74 days ago

2185cb1c

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

Google-ML-Automation committed 74 days ago

4ea536c2

Merge pull request #35696 from davisyoshida:recip-docstring

Google-ML-Automation committed 74 days ago

ea4b2e7e

davisyoshida committed 74 days ago

a767ad38

Fix reciprocal docstring

davisyoshida committed 74 days ago

f5676668

Automated Code Change

Google-ML-Automation committed 74 days ago

2d5a40fe

[ROCm] Pass s3_download_uri to pytest_rocm from orchestrator workflows

mminutoli committed 74 days ago

eb60023a

[ROCm] Wire s3_download_uri through pytest_rocm workflow

mminutoli committed 74 days ago

af6ae088

[ROCm] Download plugin wheels from S3 instead of GitHub Releases

mminutoli committed 74 days ago

595b48ce

[ROCm] Build plugin wheels for all tested Python versions

mminutoli committed 74 days ago

65331f34

[ROCm] Write LATEST pointer to S3 after wheel upload

mminutoli committed 74 days ago

51b8178a

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

yashk2810 committed 74 days ago

a364e282

[Take 2] shard map without linear util

dougalm committed 75 days ago

830e342c

Perform dtype canonicalization during Primitive.bind.

hawkinsp committed 75 days ago

fae1e236

Merge pull request #35818 from jakevdp:lax-accuracy

Google-ML-Automation committed 75 days ago

9915f8c2

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

Google-ML-Automation committed 75 days ago

964e2ebc

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

Google-ML-Automation committed 75 days ago

a8da2f48

Increase size of test jet_test_cpu to avoid timeout.

mwhittaker committed 75 days ago

326f8b44

Improve handling of accuracy parameter in jax.lax

jakevdp committed 75 days ago

fa2a8caf

Reverts 0b56550381c1f5648235606a92239c12bd1e876f

Google-ML-Automation committed 75 days ago

8c5def0f

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

Google-ML-Automation committed 75 days ago

1b268ffc

Newer Older