ROCm/jax

Pull Requests Commits

Add linear algebra tests to our CI workflow

charleshofer committed 1 year ago

94fdde79

Remove option for CI to use custom XLA (#365)

charleshofer committed 1 year ago

Verified 151f4c22

Fix backslash and quoting errors in release CI (#374)

charleshofer committed 1 year ago

Verified fda7a21e

Port Jenkins release CI to Actions (#336)

charleshofer committed 1 year ago

Verified 6451164d

Add Python and ROCm version matrix for rocm-main (#314)

charleshofer committed 1 year ago

Verified 1f93b4b9

Use rebase for upstream sync (#325)

charleshofer committed 1 year 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

Older