ROCm/jax

Pull Requests Commits

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

allanrenucci committed 75 days ago

3a971a2e

Require shape_poly_test to use at least 2 shards on TPU.

hawkinsp committed 75 days ago

21f299ef

Merge pull request #35868 from jakevdp:tracer-generic

Google-ML-Automation committed 75 days ago

b581095a

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

sohaibiftikhar committed 75 days ago

513e4bd6

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

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

21488f39

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

hawkinsp committed 75 days ago

a6db7883

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

apaszke committed 75 days ago

5be1f6fc

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

sohaibiftikhar committed 75 days ago

c4eb5b38

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

bchetioui committed 75 days ago

ce449dc0

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

bchetioui committed 75 days ago

b314b147

[Mosaic GPU] Remove the partitioned from MGPU APIs

apaszke committed 75 days ago

16798db7

[Pallas:MGPU] Remove partitioned_axis from copy_gmem_to_smem

apaszke committed 75 days ago

b5933765

Merge pull request #35880 from superbobry:pyrefly-2

Google-ML-Automation committed 75 days ago

a2c0c13a

Removed a few more unused Pyrefly ignores

superbobry committed 75 days ago

a75b3a56

[NFC] Remove obsolete TODOs in GPU layout inference.

allanrenucci committed 75 days ago

2d2f6bb0

Merge pull request #35741 from ROCm:mosaic-rocm-skips

Google-ML-Automation committed 75 days ago

2b53dd88

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

Google-ML-Automation committed 75 days ago

1ff4aeab

remove decorators, skip in setUp()

gulsumgudukbay committed 76 days ago

76341fe8

Tracer is now generic parameterized by Trace

jakevdp committed 76 days ago

8413b8a2

Merge pull request #35871 from mattjj:custom-vjp3-checkify

Google-ML-Automation committed 76 days ago

60db372e

Colocated python: Set `use_weakrefs_at_backend` to true by default.

Google-ML-Automation committed 76 days ago

0a37aaf0

[hijax][custom_vjp3] "fix" checkify to ignore custom_vjp3 decorated primals

mattjj committed 76 days ago

d8feeb68

[Mosaic TPU] Support general reshape patterns: `[src_major_dims, AxB] -> [dst_major_dims, A, B]`, where we can have arbitrary reshapes between `src_major_dims` and `dst_major_dim`.

yueshengys committed 76 days ago

19c40e08

Bugfix for colocated python with `use_weakrefs_at_backend=True`.

Google-ML-Automation committed 76 days ago

20c8683e

PR #35642: [typing] add a type signature to weakref_lru_cache

jakevdp committed 76 days ago

32f37102

Pallas: Add transpose argument to matmul_push_rhs.

jburnim committed 76 days ago

5184ae31

Expose jax.extend.core.take_current_trace

jakevdp committed 76 days ago

c294216f

[Mosaic GPU] Mark attributes of CopyPartition as class members

apaszke committed 76 days ago

07d24d89

[Pallas:MGPU] Add support for broadcasting vectors under WG semantics.

allanrenucci committed 76 days ago

cad0873a

Newer Older