ROCm/jax

Pull Requests Commits

Add `block_until_ready()` to FAQ code snippet.

Google-ML-Automation committed 299 days ago

189ba3a9

Support pltpu.roll on sublanes when not all lanes are used.

Google-ML-Automation committed 299 days ago

e4343212

[pallas:mosaic] Allowed registering lowering per `pltpu.KernelType`

superbobry committed 299 days ago

bc3a3f0b

Do all mesh checks at shard_map **call time** instead of at construction time.

yashk2810 committed 299 days ago

6fc9e17a

Default `in_axes` of `smap` to `Infer`. This matches the behavior of `jax.shard_map` where `in_specs` is optional. If the `axis_name` `smap` is going manual over is not of type `Explicit`, we error out and providing `in_axes` is compulsory.

yashk2810 committed 299 days ago

1a3d3e37

Move the existing mask handling code to the relayout fn, invoke it from the existing tpu relayout rule.

Google-ML-Automation committed 299 days ago

167d6bc7

[Mosaic] Support squeezing tiled memrefs to 1d shapes.

Google-ML-Automation committed 299 days ago

4dfcbc2c

remove :custom_call and :runtime from mosaic_gpu since they are in :mosaic_gpu_support now.

Google-ML-Automation committed 299 days ago

f6b9f7d7

Merge pull request #28695 from jakevdp:array-api-tests

Google-ML-Automation committed 299 days ago

76e5bc6f

Merge pull request #28662 from dsuedholt:faster-stft

Google-ML-Automation committed 299 days ago

b845ecdc

[array API] update test suite to most recent version

jakevdp committed 299 days ago

3b9865af

Speed up `scipy.signal.stft` by using `lax.dynamic_slice_in_dim` for windowing

davidmarttila committed 299 days ago

e65b317a

Merge pull request #28683 from Arech8:arech8/fix_bazelrc_debug

Google-ML-Automation committed 299 days ago

e2b70767

[pallas:mosaic_gpu] Slightly generalized `MosaicGridMapping`

superbobry committed 299 days ago

fc12df09

[Mosaic GPU] Add support for TMEM loads/stores with the 32x32b shape

apaszke committed 299 days ago

d99778cf

[Mosaic GPU] Use f8e4m3fn in place of f8e4m3

apaszke committed 299 days ago

f5c63053

Call block_until_ready for testAutodiffCache

Google-ML-Automation committed 299 days ago

59992087

Add collective_axes to run_scoped

apaszke committed 299 days ago

bd8765d3

Update XLA dependency to use revision

Google-ML-Automation committed 299 days ago

0c510902

[Mosaic GPU] Add an additional WG barrier before copy_gmem_to_smem

apaszke committed 299 days ago

eb2fe971

Fix debug rule in .bazelrc

Arech8 committed 299 days ago

f6bce253

[Pallas:MGPU] Update one more lowering rule to the load/store rename in TMEMRef

apaszke committed 299 days ago

863f762d

[Mosaic GPU] Add support for 8-bit MMA on Blackwell

apaszke committed 299 days ago

caf10dfc

Update XLA dependency to use revision

Google-ML-Automation committed 300 days ago

b4eb48cd

Merge pull request #28664 from mattjj:rahul-fix

Google-ML-Automation committed 301 days ago

ddfcf84a

Automated Code Change

Google-ML-Automation committed 301 days ago

2016e596

Update XLA dependency to use revision

Google-ML-Automation committed 301 days ago

6ec18d49

[si_vjp] fix bugs around symbolic zeros

mattjj committed 301 days ago

55a9de3c

[Pallas] Allow more int casting tests.

WindQAQ committed 301 days ago

35e2657b

Fix typo in FusedAttentionTest

justinjfu committed 302 days ago

0aef4475

Older