ROCm/jax

Pull Requests Commits

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

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

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

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

[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

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7a459f0e

[Mosaic GPU] Add transform inference rule for `mgpu.slice_smem`.

bchetioui committed 1 year ago

ba2f7c9a

[Mosaic GPU] Add a specialized layout for loading 4-bit inputs in WGMMA friendly layouts

apaszke committed 1 year ago

d4bd2570

[mosaic_gpu] Force flush all cupti activity, then unsubscribe.

chr1sj0nes committed 1 year ago

38d52a19

[Mosaic GPU] Remove sub-byte conversion restriction

apaszke committed 1 year ago

34cd5b0d

Older