ROCm/jax

Pull Requests Commits

[ROCm] Version-tier S3 layout, standardized CI

mminutoli committed 3 days ago

6d5fcca6

[bazel] remove `strict_deps = False` from targets that build without it.

Jake VanderPlas committed 3 days ago

aafe16fc

[bazel] support `strict_deps` argument in `py_*` rules & remove deprecated `py_strict_*` aliases.

Jake VanderPlas committed 3 days ago

eaafdff4

Add an example test for using `compute_on2` to multi-stream Pallas/Mosaic GPU kernels.

bchetioui committed 3 days ago

e4b57cd1

[pallas:sc] Updated get_p/store_p lowering rules to use full memref rank

superbobry committed 3 days ago

2176d4a8

[gh] Post-process workflow test results and load into BigQuery.

danielsuo committed 3 days ago

653c39c6

Cleanup after refactor of PJRT C API error handling utilities.

emilyfertig committed 3 days ago

dca011eb

[mosaic:gpu] Standardize test file names and targets to *_test.py.

danielsuo committed 3 days ago

ca68303c

[Pallas] Add hijax support to mpmd_map

sharadmv committed 3 days ago

79426870

Switch presubmits to use XLA at the pinned revision instead of from HEAD.

belitskiy committed 3 days ago

6b7ead69

Add a dynamic TPU Hardware Reference page inside Pallas TPU documentation.

sharadmv committed 3 days ago

8d8157f8

Merge pull request #37356 from bej9:linalg

Google-ML-Automation committed 3 days ago

6b6e7c78

[pallas] Removed one level of wrapping from `pl.kernel`

superbobry committed 3 days ago

00536f66

[Pallas/Mosaic GPU] Forbid using `barrier_test` in warpgroup contexts.

bchetioui committed 3 days ago

e9177ffd

Add a smoke test to ensure local wheels are used as expected (when expected).

belitskiy committed 3 days ago

c041fd6f

[pallas:mosaic] Inlined a few unnecessary helpers in `pltpu.emit_pipeline`

superbobry committed 3 days ago

9e8338aa

[pallas:sc] Raise instead of returning a dummy `SparseCoreInfo` in `get_sparse_core_info`

superbobry committed 3 days ago

636075ef

[Mosaic GPU] Remove duplicate word in error message.

golechwierowicz committed 3 days ago

c34ff28c

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

Google-ML-Automation committed 3 days ago

9c645a5f

[Pallas TPU] Expose number of accumulators and number of MXUs for the various TPU

bchetioui committed 3 days ago

fa71d7a8

[mgpu] Infer memory space in `ptr_as_memref` and check consistency.

chr1sj0nes committed 4 days ago

6fb32d3c

Error out in dot_general if one of lhs/rhs mesh is not empty.

Google-ML-Automation committed 4 days ago

d7f3b8f8

Add security analysis with zizmor job to ci-build.yml.

Google-ML-Automation committed 4 days ago

f05d491a

Relax type signature for io callbacks from SingleDeviceSharding -> Sharding.

pschuh committed 4 days ago

14f65de2

Error out in dot_general if one of lhs/rhs mesh is not empty.

yashk2810 committed 4 days ago

cfa823d4

Merge pull request #37405 from mattjj:linearize-instantiate-symbolic-zeros

Google-ML-Automation committed 4 days ago

39c7db58

fix forward jax.linearize symbolic zero bug from #37377

mattjj committed 4 days ago

1312fb01

Use linearize_jaxpr in compute_on linearize rule

yashk2810 committed 4 days ago

c764133f

Merge pull request #37358 from jashshah999:fix/standardize-numerical-stability

Google-ML-Automation committed 4 days ago

6335fc2f

Add algorithm parameter to nn.standardize for numerical stability

jashshah999 committed 4 days ago

6b02fbe1

Older