ROCm/jax

Pull Requests Commits

Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4

Google-ML-Automation committed 304 days ago

e9ce8fb9

Integrate LLVM at llvm/llvm-project@0230d63b4a8b

Google-ML-Automation committed 304 days ago

f3b7c5cb

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

sharadmv committed 304 days ago

e949effc

Make sure to DCE read effects

sharadmv committed 304 days ago

4d715753

Update XLA dependency to use revision

Google-ML-Automation committed 304 days ago

8c7a55ea

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

yashk2810 committed 304 days ago

663ef7ae

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

hawkinsp committed 304 days 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 304 days ago

01a110c4

Support tuples in custom_partitioning.

pschuh committed 304 days ago

0fb59747

Fix logging_test fails on Linux with NVIDIA Driver only.

Google-ML-Automation committed 304 days ago

080804c7

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

Rifur13 committed 304 days ago

54691b12

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

yashk2810 committed 304 days ago

76d9890b

fix to ragged_all_to_all transpose

mattjj committed 304 days ago

942ff38e

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

jburnim committed 304 days ago

47e8effd

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

bchetioui committed 304 days ago

875099b2

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

hawkinsp committed 304 days ago

547d6027

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

Google-ML-Automation committed 304 days 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 304 days ago

a5c0f200

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

Google-ML-Automation committed 304 days ago

7c5871f4

Merge pull request #27198 from jakevdp:lax-docs

Google-ML-Automation committed 304 days ago

30941480

Make blocked_fold_in consistent when the block sizes induce padding

Google-ML-Automation committed 304 days ago

13541e9f

jax.lax: improve docs for several APIs

jakevdp committed 304 days ago

8b46e53a

Added exit 1 if git patch is failed + other checks

vfdev-5 committed 304 days ago

9145d617

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

bchetioui committed 304 days ago

1e36cbe5

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

apaszke committed 304 days ago

8da93249

Update XLA dependency to use revision

Google-ML-Automation committed 304 days ago

7a459f0e

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

bchetioui committed 304 days ago

ba2f7c9a

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

apaszke committed 304 days ago

d4bd2570

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

chr1sj0nes committed 304 days ago

38d52a19

[Mosaic GPU] Remove sub-byte conversion restriction

apaszke committed 304 days ago

34cd5b0d

Older