ROCm/jax

Pull Requests Commits

[ROCm] Restore default SVD algorithms on ROCM.

tsrw2048 committed 103 days ago

42ef4cc1

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

chr1sj0nes committed 103 days ago

351cd419

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

allanrenucci committed 103 days ago

06669731

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

chr1sj0nes committed 103 days ago

1b6cd891

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

PatriosTheGreat committed 103 days ago

7d734b24

[Pallas:SC] Reenable a few tests

apaszke committed 103 days ago

0e3194f9

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

allanrenucci committed 103 days ago

74fa8a16

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

allanrenucci committed 103 days ago

fe5ac0fb

Merge pull request #36419 from hawkinsp:docs

Google-ML-Automation committed 103 days ago

4744f3a5

Fix dangling references to deleted documentation.

hawkinsp committed 103 days ago

eba1227a

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

allanrenucci committed 103 days ago

85549ebd

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

Google-ML-Automation committed 103 days ago

d36f6b3b

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

Google-ML-Automation committed 103 days ago

a1d03056

new-sharding-docs

mattjj committed 103 days ago

a900877d

fix typo in Pallas docs

18praveenb committed 104 days ago

220f0751

Remove xla_tpu_enable_mlir_verifier in tests.

megrez-yliu committed 104 days ago

48092bd8

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

yueshengys committed 104 days ago

ee3d8843

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

IllogicalMoose committed 104 days ago

dfc7e729

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

rdyro committed 104 days ago

dfa77ad7

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

yashk2810 committed 104 days ago

5d0b1e40

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

hawkinsp committed 104 days ago

b12e491f

PR #36375: [pyrefly] Removed unnecessary suppressions

superbobry committed 104 days ago

8480ea79

Re-land #36265 with better dtype semantics.

Jake VanderPlas committed 104 days ago

e2b1ca3b

[Pallas] Fix refactoring bug supressing TPU Interpret Mode effects

jburnim committed 104 days ago

f5fc3b9a

Skip broken SparseCore Pallas test cases.

nvgrw committed 104 days ago

faae8dd3

Document the output dtype of jax.image.resize.

hawkinsp committed 104 days ago

798e65b1

[Mosaic GPU][NFC] Remove unnecessary lambdas.

allanrenucci committed 104 days ago

8f40c043

Reverts 6e5a5ca7427af8a6d4c732065f42abb0edca7574

hawkinsp committed 104 days ago

5ad89659

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

bchetioui committed 104 days ago

017dd5d1

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

hawkinsp committed 104 days ago

8dc3715a

Older