ROCm/jax

Pull Requests Commits

Add canonicalization of python versions

mrodden committed 296 days ago

5bf08753

Split parse_gpu_targets test cases

mrodden committed 296 days ago

5024cf8a

Fix missing defaults for GPU target arch'es (#293)

mrodden committed 296 days ago

ba2e1128

Use rebase for upstream sync (#325)

charleshofer committed 297 days ago

Verified 7b20b7c0

Fixes for 0.5.0 build ported to rocm-main

mrodden committed 299 days ago

ec4b8ee1

Add back raw totals in JSON reports (#281)

charleshofer committed 300 days ago

Verified 13d88b63

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

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

Verified b505df99

Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4

Google-ML-Automation committed 306 days ago

e9ce8fb9

Integrate LLVM at llvm/llvm-project@0230d63b4a8b

Google-ML-Automation committed 306 days ago

f3b7c5cb

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

sharadmv committed 306 days ago

e949effc

Make sure to DCE read effects

sharadmv committed 306 days ago

4d715753

Update XLA dependency to use revision

Google-ML-Automation committed 306 days ago

8c7a55ea

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

yashk2810 committed 306 days ago

663ef7ae

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

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

01a110c4

Enable upstream CI on release branches (#296)

charleshofer committed 306 days ago

Verified 1f2fe330

Support tuples in custom_partitioning.

pschuh committed 306 days ago

0fb59747

Fix logging_test fails on Linux with NVIDIA Driver only.

Google-ML-Automation committed 306 days ago

080804c7

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

Rifur13 committed 306 days ago

54691b12

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

yashk2810 committed 306 days ago

76d9890b

fix to ragged_all_to_all transpose

mattjj committed 306 days ago

942ff38e

Fix ROCm build README (#284)

charleshofer committed 306 days ago

Verified dd7f96b2

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

jburnim committed 306 days ago

47e8effd

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

bchetioui committed 306 days ago

875099b2

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

hawkinsp committed 306 days ago

547d6027

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

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

a5c0f200

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

Google-ML-Automation committed 306 days ago

7c5871f4

Merge pull request #27198 from jakevdp:lax-docs

Google-ML-Automation committed 306 days ago

30941480

Make blocked_fold_in consistent when the block sizes induce padding

Google-ML-Automation committed 306 days ago

13541e9f

Older