ROCm/jax

Pull Requests Commits

Merge pull request #35696 from davisyoshida:recip-docstring

Google-ML-Automation committed 121 days ago

ea4b2e7e

davisyoshida committed 121 days ago

a767ad38

Fix reciprocal docstring

davisyoshida committed 121 days ago

f5676668

Automated Code Change

Google-ML-Automation committed 121 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 122 days ago

a364e282

[Take 2] shard map without linear util

dougalm committed 122 days ago

830e342c

Perform dtype canonicalization during Primitive.bind.

hawkinsp committed 122 days ago

fae1e236

Merge pull request #35818 from jakevdp:lax-accuracy

Google-ML-Automation committed 122 days ago

9915f8c2

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

Google-ML-Automation committed 122 days ago

964e2ebc

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

Google-ML-Automation committed 122 days ago

a8da2f48

Increase size of test jet_test_cpu to avoid timeout.

mwhittaker committed 122 days ago

326f8b44

Improve handling of accuracy parameter in jax.lax

jakevdp committed 122 days ago

fa2a8caf

Reverts 0b56550381c1f5648235606a92239c12bd1e876f

Google-ML-Automation committed 122 days ago

8c5def0f

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

Google-ML-Automation committed 122 days ago

1b268ffc

Merge pull request #35894 from hawkinsp:docports

Google-ML-Automation committed 122 days ago

573b73d7

Attempt 2 at fixing doc build port collisions.

hawkinsp committed 122 days ago

3e463e76

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

gulsumgudukbay committed 122 days ago

2fb2d6cf

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

apaszke committed 122 days ago

4acb04de

[Pallas:MGPU] Add Pallas lowering for atomic stores under WG semantic.

allanrenucci committed 122 days ago

3a971a2e

Require shape_poly_test to use at least 2 shards on TPU.

hawkinsp committed 122 days ago

21f299ef

Merge pull request #35868 from jakevdp:tracer-generic

Google-ML-Automation committed 122 days ago

b581095a

[MOSAIC:GPU]: Align collective kernel metadata size to XLA

sohaibiftikhar committed 122 days ago

513e4bd6

Change TypedNdArray to be a subclass of np.ndarray, rather than a duck type.

hawkinsp committed 122 days ago

dda96957

Fix a bug in sort_key_val for negative integers, since SC sort treats as unsigned. Also adds support for u32. (Similar to an earlier fix for reduce min/max.)

brianwa84 committed 122 days ago

21488f39

Use a NumPy reference implementation in tests/pallas/einshape_test.py::EinshapeTest::test_hypothesis_einshape.

hawkinsp committed 122 days ago

a6db7883

[Mosaic GPU] Add support for s4 warp-level MMAs

apaszke committed 122 days ago

5be1f6fc

[XLA:GPU]: Align collective kernel metadata to 16 bytes

sohaibiftikhar committed 122 days ago

c4eb5b38

[Pallas/Mosaic GPU] Fix misuse of `self.{kernel,pallas_call}` in `mosaic_gpu_test.py`.

bchetioui committed 122 days ago

ce449dc0

[Pallas/Mosaic GPU] Turn on `gpu_pallas_distributed_test`s for warpgroup semantics.

bchetioui committed 122 days ago

b314b147

[Mosaic GPU] Remove the partitioned from MGPU APIs

apaszke committed 122 days ago

16798db7

Newer Older