ROCm/jax

Pull Requests Commits

[pallas:triton] Fix atomic min/max lowering for unsigned integers and float types

draganmladjenovic committed 303 days ago

5a59b1dd

Use consistent dtype for forward and backwards in jax.nn.dot_product_attention.

sbodenstein committed 311 days ago

d5e5b42d

Update XLA dependency to use revision

Google-ML-Automation committed 311 days ago

44872296

Expose `get_ty` aka get_aval from jax namespace

yashk2810 committed 311 days ago

b3508333

Exclude auto axes whenever extending the axis_env via core.extend_axis_env_nd.

pschuh committed 311 days ago

c825241c

Merge pull request #26172 from ZacCranko:is-distributed-init

Google-ML-Automation committed 311 days ago

09491e2b

add distributed.is_initialized

ZacCranko committed 311 days ago

5db78e7a

Let users pass in pspecs to with_sharding_constraint when `use_mesh` is set. This is in-line with other APIs which allow pspecs like einsum, reshape, etc

yashk2810 committed 311 days ago

1079dc44

Make device_put resharding on single device array input work under use_mesh. Fixes https://github.com/jax-ml/jax/issues/26552

yashk2810 committed 311 days ago

8bcbf585

[sharding_in_types] Set the `sharding_in_types` config to True. This is a purely internal change and shouldn't affect any public APIs.

yashk2810 committed 311 days ago

00d82970

[Mosaic TPU] Support mask concat

bythew3i committed 311 days ago

bb68124c

Error unconditionally for jit, pjit and with_sharding_constraint if `use_mesh` and `with mesh` are used together.

yashk2810 committed 311 days ago

1dc58b79

Merge pull request #26401 from jakevdp:numpy-consts

Google-ML-Automation committed 311 days ago

72f0a90e

[pallas] Skip `OpsTest.test_concat_constant` if TPU is not available

superbobry committed 311 days ago

d4559ba4

Integrate LLVM at llvm/llvm-project@9d24f9437944

Google-ML-Automation committed 311 days ago

725087e1

Merge pull request #26574 from BaconBreaker:lexsort-stub-file

Google-ML-Automation committed 311 days ago

82216fc4

Merge pull request #26550 from jakevdp:lax-docs

Google-ML-Automation committed 311 days ago

63bc22d6

Merge pull request #26571 from rajasekharporeddy:typos

Google-ML-Automation committed 311 days ago

8f8e3187

Update XLA dependency to use revision

Google-ML-Automation committed 311 days ago

ce399aa6

Update XLA dependency to use revision

Google-ML-Automation committed 311 days ago

9f8a0bf9

Update XLA dependency to use revision

Google-ML-Automation committed 312 days ago

490ceba7

Merge pull request #26570 from hawkinsp:pgle

Google-ML-Automation committed 312 days ago

7e27713a

Change type signature of lexsort in stub file to match type signature in sorting.py

BaconBreaker committed 312 days ago

422b747d

rajasekharporeddy committed 312 days ago

180be997

Fix test failure in PGLE tests.

hawkinsp committed 312 days ago

b3ed528f

Integrate LLVM at llvm/llvm-project@912b154f3a3f

Google-ML-Automation committed 312 days ago

e78a469b

Update XLA dependency to use revision

Google-ML-Automation committed 313 days ago

bffeb5d6

[Mosaic GPU] Implement lowerings for `Tile` and `Transpose` transforms from the MLIR dialect.

dimitar-asenov committed 313 days ago

52f8fbee

Cache `shape` property on `Mesh`.

bartchr808 committed 313 days ago

7bfa420d

[TPU][Mosaic][Easy] Add verification for AssumeMultipleOp.

Google-ML-Automation committed 313 days ago

a6fcb741

Older