ROCm/jax

Pull Requests Commits

[Pallas:MGPU] Remove partitioned_axis from copy_gmem_to_smem

apaszke committed 122 days ago

b5933765

Merge pull request #35880 from superbobry:pyrefly-2

Google-ML-Automation committed 122 days ago

a2c0c13a

Removed a few more unused Pyrefly ignores

superbobry committed 122 days ago

a75b3a56

[NFC] Remove obsolete TODOs in GPU layout inference.

allanrenucci committed 122 days ago

2d2f6bb0

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

Google-ML-Automation committed 122 days ago

2b53dd88

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

Google-ML-Automation committed 122 days ago

1ff4aeab

remove decorators, skip in setUp()

gulsumgudukbay committed 123 days ago

76341fe8

Tracer is now generic parameterized by Trace

jakevdp committed 123 days ago

8413b8a2

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

Google-ML-Automation committed 123 days ago

60db372e

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

Google-ML-Automation committed 123 days ago

0a37aaf0

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

mattjj committed 123 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 123 days ago

19c40e08

Bugfix for colocated python with `use_weakrefs_at_backend=True`.

Google-ML-Automation committed 123 days ago

20c8683e

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

jakevdp committed 123 days ago

32f37102

Pallas: Add transpose argument to matmul_push_rhs.

jburnim committed 123 days ago

5184ae31

Expose jax.extend.core.take_current_trace

jakevdp committed 123 days ago

c294216f

[Mosaic GPU] Mark attributes of CopyPartition as class members

apaszke committed 123 days ago

07d24d89

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

allanrenucci committed 123 days ago

cad0873a

Don't eagerly compute error messages, especially in assertAllClose.

hawkinsp committed 123 days ago

5ff3a9e8

[Mosaic GPU] Add layout inference rule to AsymcStoreSparseMetadataSmemToTmem op.

golechwierowicz committed 123 days ago

00c40ca8

Apply the documentation fix from https://github.com/jax-ml/jax/pull/20828 using math notation.

hawkinsp committed 123 days ago

f06655b9

[Mosaic GPU] Add AsyncStoreSparseMetadataSmemToTmemOp op.

golechwierowicz committed 123 days ago

69996b16

Apply a fix from PR https://github.com/jax-ml/jax/pull/21020.

hawkinsp committed 123 days ago

24f387da

Merge pull request #35855 from jakevdp:array-overrides

Google-ML-Automation committed 123 days ago

7303b4e1

Remove jax/tools/colab_tpu.py.

hawkinsp committed 123 days ago

67da5f94

Fix an inappropriate use of /tmp in a test on Windows.

hawkinsp committed 123 days ago

9d8d5bc5

Merge pull request #35856 from superbobry:pyrefly-2

Google-ML-Automation committed 123 days ago

41d3f9fd

Merge pull request #35841 from gnecula:jax2tf_named_sharding

Google-ML-Automation committed 123 days ago

726c2244

[mlir] Fixed a few type errors in MLIR usages

superbobry committed 123 days ago

e67f9dc3

Change unreduced/reduced rule's signature to only take avals and params. Also combine unreduced and reduced rule because right now they have a very finicky reduced -> unreduced rule dependency.

yashk2810 committed 123 days ago

f94b79d9

Newer Older