ROCm/jax

Pull Requests Commits

Add linear algebra tests to our CI workflow

charleshofer committed 237 days ago

94fdde79

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

charleshofer committed 242 days ago

Verified 151f4c22

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

charleshofer committed 242 days ago

Verified fda7a21e

Port Jenkins release CI to Actions (#336)

charleshofer committed 249 days ago

Verified 6451164d

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

charleshofer committed 262 days ago

Verified 1f93b4b9

Use rebase for upstream sync (#325)

charleshofer committed 268 days ago

Verified 7b20b7c0

Fixes for 0.5.0 build ported to rocm-main

mrodden committed 270 days ago

ec4b8ee1

Add back raw totals in JSON reports (#281)

charleshofer committed 271 days ago

Verified 13d88b63

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

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

Verified b505df99

Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4

Google-ML-Automation committed 277 days ago

e9ce8fb9

Integrate LLVM at llvm/llvm-project@0230d63b4a8b

Google-ML-Automation committed 277 days ago

f3b7c5cb

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

sharadmv committed 277 days ago

e949effc

Make sure to DCE read effects

sharadmv committed 277 days ago

4d715753

Update XLA dependency to use revision

Google-ML-Automation committed 277 days ago

8c7a55ea

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

yashk2810 committed 277 days ago

663ef7ae

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

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

01a110c4

Enable upstream CI on release branches (#296)

charleshofer committed 277 days ago

Verified 1f2fe330

Support tuples in custom_partitioning.

pschuh committed 277 days ago

0fb59747

Fix logging_test fails on Linux with NVIDIA Driver only.

Google-ML-Automation committed 277 days ago

080804c7

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

Rifur13 committed 277 days ago

54691b12

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

yashk2810 committed 277 days ago

76d9890b

fix to ragged_all_to_all transpose

mattjj committed 277 days ago

942ff38e

Fix ROCm build README (#284)

charleshofer committed 277 days ago

Verified dd7f96b2

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

jburnim committed 277 days ago

47e8effd

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

bchetioui committed 277 days ago

875099b2

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

hawkinsp committed 277 days ago

547d6027

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

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

a5c0f200

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

Google-ML-Automation committed 277 days ago

7c5871f4

Older