ROCm/jax

Pull Requests Commits

Merge pull request #299 from ROCm/ci-upstream-sync-152_1

rocm-repo-management-api-2[bot] committed 1 year ago

Verified b505df99

Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4

Google-ML-Automation committed 1 year ago

e9ce8fb9

Integrate LLVM at llvm/llvm-project@0230d63b4a8b

Google-ML-Automation committed 1 year ago

f3b7c5cb

[Pallas/Fuser] DCE fusion jaxprs before pulling (to avoid unnecessary computations being staged out in block functions)

sharadmv committed 1 year ago

e949effc

Make sure to DCE read effects

sharadmv committed 1 year ago

4d715753

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

8c7a55ea

Check the type of mesh in `use_abstract_mesh` and `use_concrete_mesh`

yashk2810 committed 1 year ago

663ef7ae

Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/

hawkinsp committed 1 year ago

3f91b4b4

Better mosaic lowering for dynamic shapes, extend an interpreter into shape_poly dimexpr and lower them alongside the graph if we are in a dynamic export regime.

Google-ML-Automation committed 1 year ago

01a110c4

Enable upstream CI on release branches (#296)

charleshofer committed 1 year ago

Verified 1f2fe330

Support tuples in custom_partitioning.

pschuh committed 1 year ago

0fb59747

Fix logging_test fails on Linux with NVIDIA Driver only.

Google-ML-Automation committed 1 year ago

080804c7

[Mosaic GPU] Support reads/writes from SMEM to WGMMARowFragLayout arrays.

Rifur13 committed 1 year ago

54691b12

Run the stream annotation tests on 2 devices so that it can be tested in TAP

yashk2810 committed 1 year ago

76d9890b

fix to ragged_all_to_all transpose

mattjj committed 1 year ago

942ff38e

Fix ROCm build README (#284)

charleshofer committed 1 year ago

Verified dd7f96b2

Adds option to initialize buffers to NaNs or zeros in TPU interpret mode.

jburnim committed 1 year ago

47e8effd

[Mosaic GPU] Enable the new transform inference pass in the warpgroup lowering.

bchetioui committed 1 year ago

875099b2

Remove //jaxlib:cpu_kernels and //jaxlib:gpu_kernels forwarding Bazel targets.

hawkinsp committed 1 year ago

547d6027

Merge pull request #27094 from vfdev-5:fix-tsan-numpy-install-patch

Google-ML-Automation committed 1 year ago

ee0073e6

`set_mesh` should return the prev_mesh instead of nothing. Users can choose to use the return value or ignore it.

yashk2810 committed 1 year ago

a5c0f200

[Pallas TPU] Hoist prologue and epilogue outside of pipeline loop

Google-ML-Automation committed 1 year ago

7c5871f4

Merge pull request #27198 from jakevdp:lax-docs

Google-ML-Automation committed 1 year ago

30941480

Make blocked_fold_in consistent when the block sizes induce padding

Google-ML-Automation committed 1 year ago

13541e9f

Merge pull request #294 from ROCm/ci-upstream-sync-151_1

charleshofer committed 1 year ago

Verified c46b4fc0

jax.lax: improve docs for several APIs

jakevdp committed 1 year ago

8b46e53a

Added exit 1 if git patch is failed + other checks

vfdev-5 committed 1 year ago

9145d617

Merge branch 'rocm-main' into ci-upstream-sync-151_1

charleshofer committed 1 year ago

c7b407c9

[Mosaic GPU] Raise a `NotImplementedError` if `swizzle=16`.

bchetioui committed 1 year ago

1e36cbe5

[Mosaic GPU] Fuse slicing into s4 -> bf16 upcasts

apaszke committed 1 year ago

8da93249

Older