ROCm/jax

Pull Requests Commits

update to fix lint for pre-commit run

cj401-amd committed 14 days ago

e8d3c8c7

update with a simplified tests for rocm profiling

cj401-amd committed 14 days ago

bbe61c27

re-format with 2-space indent

cj401-amd committed 14 days ago

4b533b20

update test for rocm profiling

cj401-amd committed 14 days ago

9bdcb8a5

update unit test with kernel_details

cj401-amd committed 14 days ago

febc59fe

Address Gemini review

JehandadKhan committed 14 days ago

25ba64c0

[ROCm] update to test if there are GPU events when doing profiling on ROCm via matmul

cj401-amd committed 14 days ago

3f8df9d2

Merge pull request #35082 from jakevdp:fix-jupytext

Google-ML-Automation committed 14 days ago

a18a7648

Fix jupytext lint error

jakevdp committed 14 days ago

77987113

Removed more dead code all over JAX

superbobry committed 14 days ago

b38a8164

Merge pull request #35047 from mattjj:shmap-docs-again

Google-ML-Automation committed 14 days ago

019c6d63

[pallas:mosaic_gpu] Removed special-casing for `kernel_name` in `plgpu.kernel`

superbobry committed 14 days ago

bfaaac9b

[mosaic_gpu] Removed a few unused functions in `inference_utils`

superbobry committed 14 days ago

a6aab034

enable shmap docs build (again)

mattjj committed 14 days ago

8181e501

Merge pull request #35052 from jakevdp:map-bugs

Google-ML-Automation committed 14 days ago

2b2d52df

Update `rules_ml_toolchain` to pass CCCL headers to NVCC before all other headers for local CUDA.

ybaturina committed 14 days ago

04c74b37

[pallas] Updated CHANGELOG.md

superbobry committed 14 days ago

b6c489d9

Request a full h100 gpu when neccessary.

Google-ML-Automation committed 14 days ago

95ab165e

Merge pull request #35051 from jakevdp:pyrefly-fixes-4

Google-ML-Automation committed 14 days ago

fcf0aa55

[XLA:MGPU] Remove Tiling and Replicated pythonic impl.

golechwierowicz committed 14 days ago

9705c50a

[XLA:MGPU] Call init_cc_mlir in jaxlib __init__.py

golechwierowicz committed 14 days ago

979067ea

[pyrefly] more type annotation fixes

jakevdp committed 14 days ago

9b258b8a

[pyrefly] add suppressions for pyrefly bug related to redefined builtins

jakevdp committed 14 days ago

d3bbb319

[XLA:MGPU] Port transfer plan classes to C++.

golechwierowicz committed 14 days ago

ba95b811

Merge pull request #35043 from jakevdp:pyrefly-iinfo

Google-ML-Automation committed 14 days ago

2dc85d45

Removed unused `str_eqn_compact`

superbobry committed 14 days ago

c45e0bf8

Merge pull request #34870 from ROCm:add-rocm-lu-pivots-to-permutation-test

Google-ML-Automation committed 14 days ago

32d2f29c

Update XLA dependency to use revision http://github.com/openxla/xla/commit/e0a75749fa93ecff45e6b97243902addb72cd7b2

Google-ML-Automation committed 14 days ago

67c21e1c

Merge pull request #34501 from magaonka-amd:enable-cuda-array-interface-upstream

Google-ML-Automation committed 14 days ago

1429df92

Merge pull request #34894 from ROCm:rocm-cholesky-solver-potrf

Google-ML-Automation committed 14 days ago

461096ea

Older