ROCm/jax

Pull Requests Commits

Add canonicalization of python versions

mrodden committed 362 days ago

5bf08753

Split parse_gpu_targets test cases

mrodden committed 362 days ago

5024cf8a

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

mrodden committed 362 days ago

ba2e1128

Use rebase for upstream sync (#325)

charleshofer committed 363 days ago

Verified 7b20b7c0

Fixes for 0.5.0 build ported to rocm-main

mrodden committed 1 year ago

ec4b8ee1

Add back raw totals in JSON reports (#281)

charleshofer committed 1 year ago

Verified 13d88b63

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

Older