ROCm/jax

Pull Requests Commits

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

rocm-repo-management-api-2[bot] committed 283 days ago

Verified b505df99

Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4

Google-ML-Automation committed 283 days ago

e9ce8fb9

Integrate LLVM at llvm/llvm-project@0230d63b4a8b

Google-ML-Automation committed 283 days ago

f3b7c5cb

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

sharadmv committed 283 days ago

e949effc

Make sure to DCE read effects

sharadmv committed 283 days ago

4d715753

Update XLA dependency to use revision

Google-ML-Automation committed 283 days ago

8c7a55ea

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

yashk2810 committed 283 days ago

663ef7ae

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

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

01a110c4

Enable upstream CI on release branches (#296)

charleshofer committed 283 days ago

Verified 1f2fe330

Support tuples in custom_partitioning.

pschuh committed 284 days ago

0fb59747

Fix logging_test fails on Linux with NVIDIA Driver only.

Google-ML-Automation committed 284 days ago

080804c7

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

Rifur13 committed 284 days ago

54691b12

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

yashk2810 committed 284 days ago

76d9890b

fix to ragged_all_to_all transpose

mattjj committed 284 days ago

942ff38e

Fix ROCm build README (#284)

charleshofer committed 284 days ago

Verified dd7f96b2

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

jburnim committed 284 days ago

47e8effd

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

bchetioui committed 284 days ago

875099b2

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

hawkinsp committed 284 days ago

547d6027

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

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

a5c0f200

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

Google-ML-Automation committed 284 days ago

7c5871f4

Merge pull request #27198 from jakevdp:lax-docs

Google-ML-Automation committed 284 days ago

30941480

Make blocked_fold_in consistent when the block sizes induce padding

Google-ML-Automation committed 284 days ago

13541e9f

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

charleshofer committed 284 days ago

Verified c46b4fc0

jax.lax: improve docs for several APIs

jakevdp committed 284 days ago

8b46e53a

Added exit 1 if git patch is failed + other checks

vfdev-5 committed 284 days ago

9145d617

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

charleshofer committed 284 days ago

c7b407c9

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

bchetioui committed 284 days ago

1e36cbe5

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

apaszke committed 284 days ago

8da93249

Older