ROCm/jax

Pull Requests Commits

[ROCm] Limiting Pallas kerenels tiles to 64 for rocm platforms

zahiqbal committed 1 year ago

93f9c2e9

[Mosaic TPU] Fold sublane offset to indices when storing to untiled ref.

bythew3i committed 1 year ago

f899d515

Remove unneeded dependency from rocm_plugin_extension.

klucke committed 1 year ago

f3e7e682

Merge pull request #25009 from jakevdp:keyarray

Google-ML-Automation committed 1 year ago

c1ae13b0

[pallas:mosaic_gpu] `emit_pipeline` now correctly supports `BlockSpec`s in GMEM

superbobry committed 1 year ago

1efef6bf

Fix false positive `debug_nans` error caused by NaNs that are properly handled in `jax.scipy.stats.gamma`

dfm committed 1 year ago

96c01299

Merge pull request #25034 from gnecula:poly_state

Google-ML-Automation committed 1 year ago

e707edea

[pallas] Add more test cases for Triton bitcast_convert_type lowering rule.

petebu committed 1 year ago

2178ed2f

[mgpu] Pointwise op can handle LHS splats.

cperivol committed 1 year ago

1d2dc17e

Merge pull request #25018 from jakevdp:update-array-api

Google-ML-Automation committed 1 year ago

b1b1ad62

Fix cron schedule to run past minute 0 every 2nd hour

nitins17 committed 1 year ago

1e6654a0

Merge pull request #25015 from barnesjoseph:add-google-sans

Google-ML-Automation committed 1 year ago

73352677

[JAX] Ignore xla_gpu_experimental_autotune_cache_mode when calculating module hash.

Google-ML-Automation committed 1 year ago

bf0150bb

Run the TPU workflow on new self-hosted runners

nitins17 committed 1 year ago

7d7a0fa2

[shape_poly] Adding shape polymorphism support for the state primitives.

gnecula committed 1 year ago

0831e2e3

Integrate LLVM at llvm/llvm-project@33fcd6acc755

metaflow committed 1 year ago

1bc9df42

[pallas:mosaic_gpu] Pulled `delay_release` into `emit_pipeline`

superbobry committed 1 year ago

f18df8f3

Reverts c04aec9d525dd2e767495e41b98e82dd79315f37

naummo committed 1 year ago

e72b4490

[sharding_in_types] Add `concatenate_p` support

yashk2810 committed 1 year ago

6568713a

[Mosaic TPU] Add bound check for general vector store op.

bythew3i committed 1 year ago

869a5334

[sharding_in_types] Add `pad_p` support to sharding_in_types to handle transpose to slice correctly.

yashk2810 committed 1 year ago

840cf3f7

[Pallas] Use Pallas cost estimator for flash attention.

justinjfu committed 1 year ago

1f6152d1

Merge pull request #25020 from jakevdp:lax-pad-validation

Google-ML-Automation committed 1 year ago

f39392ea

Adds Google Sans font

barnesjoseph committed 1 year ago

bf7f9aa8

jax.lax.pad: improve input validation

jakevdp committed 1 year ago

17825882

Merge pull request #25019 from jakevdp:lax-pad-doc

Google-ML-Automation committed 1 year ago

334bd4d0

DOC: add examples for jax.lax.pad

jakevdp committed 1 year ago

2699e950

[array api] use most recent version of array_api_tests

jakevdp committed 1 year ago

f749fca7

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6fe78042

[sharding_in_types] Add slice_p and squeeze_p sharding rule to make flash attention work in backward pass

yashk2810 committed 1 year ago

9b941808

Older