ROCm/jax

Pull Requests Commits

lax.top_k: add axis support

jakevdp committed 185 days ago

990c366f

Add sharding checks in `_primal_tangent_shapes_match` that is called in the constructor of `LinearizeTracer`.

yashk2810 committed 185 days ago

30582db2

[pallas:sc] Added additional checks for scatter/gather via `pltpu.async_copy`

superbobry committed 185 days ago

147f6077

[Pallas:MGPU] Reject transposed TMEM references.

allanrenucci committed 185 days ago

1a01c082

[Pallas:MGPU][NFC] Simplify test using `functools.partial`.

allanrenucci committed 185 days ago

0195d356

[Mosaic:TPU][apply-vector-layout] tpu.concatenate for result offsets != 0

tlongeri committed 185 days ago

f912fd71

[Pallas:MGPU] Add warpgroup lowering for `tcgen05_commit_arrive`.

allanrenucci committed 185 days ago

e00f9c03

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

Google-ML-Automation committed 185 days ago

21291ca4

Add a test which shows that fsdp matmul with unreduced works without custom_vjp

yashk2810 committed 186 days ago

7bdcd36c

kill em with kindness

mattjj committed 186 days ago

65401738

Exposes some hijax APIs required to develop custom hijax types in jax.experimental.hijax

cgarciae committed 186 days ago

cce23c06

[Mosaic TPU] Add pack/unpack elementwise operations.

megrez-yliu committed 186 days ago

af5b5c8a

dougalm committed 186 days ago

2e6f1951

[Mosaic TPU] Support more general i1 relayout which usually comes from masking. Before this change, we will pick too small sublane tiling for packed types during relayout insertion, e.g., choose (1, 128) for i8.

yueshengys committed 186 days ago

7717ae41

dougalm committed 186 days ago

60bcd6af

Merge pull request #33059 from yliu120:patch-3

Google-ML-Automation committed 186 days ago

6133d7e2

another attempt: hogwild refs are different

mattjj committed 186 days ago

54605f23

Merge pull request #31469 from yliu120:yl/deserialize_exec

Google-ML-Automation committed 186 days ago

e9609ce4

Add `out_sharding` parameter to `_IndexUpdateRef` `get/set/add` methods in .pyi files

yashk2810 committed 186 days ago

2f8838a7

Fix handling of pytrees within `jnp.array` and `jnp.asarray`.

Jake VanderPlas committed 186 days ago

23c7ae9b

Reverts 66924519831dddb9ede68cd868052a708a3d4e15

hawkinsp committed 186 days ago

2a8cb54b

Adds deserializing exec capability with new FfiHostCallback API.

yliu120 committed 186 days ago

c10ffa32

[Mosaic TPU] Add a missing check in load->reshape optimization

apaszke committed 186 days ago

ce3dd1fe

Make jax.grad work when the output on the fwd pass is unreduced.

yashk2810 committed 186 days ago

8050eddd

7x device mesh fall back to unoptimized when needed

samos123 committed 186 days ago

75b11ec2

Add partial_eval custom rule for xla_metadata_call_p so that it works under a remat

yashk2810 committed 186 days ago

44f71433

[Mosaic TPU] Fix up some broken tests

apaszke committed 186 days ago

8777357b

[Pallas:MGPU] Add initial warpgroup lowering for `tcgen05_mma`.

allanrenucci committed 186 days ago

9bad30c1

Prepare jaxlib for ROCm 7.0

beckerhe committed 186 days ago

090f73de

[Pallas:MGPU][NFC] Remove duplicate logic to handle transposed refs in WGMMA WG lowering.

allanrenucci committed 186 days ago

0f63a82a

Newer Older