ROCm/jax

Pull Requests Commits

[ROCm] Change run_multi_gpu set opts

Ruturaj4 committed 1 year ago

694de6b6

[sharding_in_types] Add scan support to sharding_in_types. There are a couple of changes here

yashk2810 committed 1 year ago

355589f3

Merge pull request #25048 from jax-ml:linearization-rule-signature

Google-ML-Automation committed 1 year ago

3d79df24

Change signature of linearization rules.

dougalm committed 1 year ago

170718c8

[Pallas] Add readme page for debugging tips.

justinjfu committed 1 year ago

344d0d99

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

26443bbd

[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

Older