ROCm/jax

Pull Requests Commits

[ROCm] Restore default SVD algorithms on ROCM.

tsrw2048 committed 5 hours ago

42ef4cc1

[pallas:mgpu] Support `Element` block dim in pipeline block specs.

chr1sj0nes committed 6 hours ago

351cd419

[Pallas:MGPU][NFC] Update docstring for `tcgen05_commit_arrive`.

allanrenucci committed 6 hours ago

06669731

[pallas:triton] Add check for cast to/from `float8_e4m3fn`.

chr1sj0nes committed 6 hours ago

1b6cd891

[Mosaic:GPU] Use XLA_VLOG_DEVICE in Mosaic GPU custom call.

PatriosTheGreat committed 6 hours ago

7d734b24

[Pallas:SC] Reenable a few tests

apaszke committed 7 hours ago

0e3194f9

[Mosaic GPU][NFC] Remove redundant `static_cast`.

allanrenucci committed 7 hours ago

74fa8a16

[Pallas:MGPU] Add support for `tcgen05_commit_arrive` under WGxWARP semantics.

allanrenucci committed 7 hours ago

fe5ac0fb

Merge pull request #36419 from hawkinsp:docs

Google-ML-Automation committed 7 hours ago

4744f3a5

Fix dangling references to deleted documentation.

hawkinsp committed 7 hours ago

eba1227a

[Mosaic GPU][NFC] Simplify `TileTransformAttr` and `TransposeTransformAttr` Python bindings.

allanrenucci committed 8 hours ago

85549ebd

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

Google-ML-Automation committed 11 hours ago

d36f6b3b

Merge pull request #36402 from jax-ml:new-sharding-docs

Google-ML-Automation committed 19 hours ago

a1d03056

new-sharding-docs

mattjj committed 20 hours ago

a900877d

fix typo in Pallas docs

18praveenb committed 21 hours ago

220f0751

Remove xla_tpu_enable_mlir_verifier in tests.

megrez-yliu committed 22 hours ago

48092bd8

[Mosaic TPU] Support 16-bit `arith::CmpI` and 1d tiling subelement mask on TPU v4 and v5.

yueshengys committed 22 hours ago

ee3d8843

Replace `use_magma=True` with `implementation=lax_linalg.EigImplementation.MAGMA` in magma_linalg_test.py.

IllogicalMoose committed 22 hours ago

dfc7e729

Finally remove the `vma` now that `manual_axis_type.varying` is the default.

rdyro committed 22 hours ago

dfa77ad7

Fix `mat.invarying` by looking at manual axes of mesh instead of all mesh axes.

yashk2810 committed 23 hours ago

5d0b1e40

Validate that an explicit axis_size passed to vmap() matches the sizes of the mapped axes.

hawkinsp committed 23 hours ago

b12e491f

PR #36375: [pyrefly] Removed unnecessary suppressions

superbobry committed 23 hours ago

8480ea79

Re-land #36265 with better dtype semantics.

Jake VanderPlas committed 23 hours ago

e2b1ca3b

[Pallas] Fix refactoring bug supressing TPU Interpret Mode effects

jburnim committed 1 day ago

f5fc3b9a

Skip broken SparseCore Pallas test cases.

nvgrw committed 1 day ago

faae8dd3

Document the output dtype of jax.image.resize.

hawkinsp committed 1 day ago

798e65b1

[Mosaic GPU][NFC] Remove unnecessary lambdas.

allanrenucci committed 1 day ago

8f40c043

Reverts 6e5a5ca7427af8a6d4c732065f42abb0edca7574

hawkinsp committed 1 day ago

5ad89659

[Mosaic GPU] Support propagating arbitrary tilings without associated swizzle.

bchetioui committed 1 day ago

017dd5d1

Fall back to the pure JAX implementation of tridiagonal_solve for m<=2 on GPU:

hawkinsp committed 1 day ago

8dc3715a

Older