ROCm/jax

Pull Requests Commits

Migrate export_test to use jit instead of pjit

gnecula committed 326 days ago

3d6fa3ba

[pallas:mosaic] Fixed the `NDIndexer` instantiation in `_prng_key_load_lowering_rule`

superbobry committed 326 days ago

3ebe0181

[Pallas:MGPU] Support reduce_scatter along arbitrary tensor dimensions

apaszke committed 326 days ago

c0165217

[Pallas:MGPU] Try to collapse dimension pairs in transfer_strided

apaszke committed 326 days ago

ce52ce80

[Pallas:MGPU] Add a small benchmark script to the reduce scatter kernel

apaszke committed 326 days ago

32add9e6

Merge pull request #32514 from apaszke:symm2

Google-ML-Automation committed 326 days ago

300431da

[Mosaic GPU] Add an SMEM inference rule for `memref.transpose`.

dimitar-asenov committed 326 days ago

414f039d

Use `()` instead of `tuple()` and `{}` instead of `dict()`

superbobry committed 326 days ago

f0540d11

[Mosaic GPU] Fix error message.

dimitar-asenov committed 326 days ago

70d266ee

[Mosaic GPU] Add SMEM inference rule for `mgpu.custom_op`.

dimitar-asenov committed 326 days ago

260d3f0f

[Mosaic GPU] Add a basic test for running collectives on Torch tensors

apaszke committed 326 days ago

524edf29

[Mosaic GPU] Add support for collective async copies.

allanrenucci committed 326 days ago

b4e24590

[Mosaic GPU][Serde] Add upgrade rules for vector.splat.

Google-ML-Automation committed 326 days ago

6c97ee88

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

Google-ML-Automation committed 326 days ago

4e14ea10

Reverts 18f96c9009624d6bec0cade1f514b8f2601f4d6a

pravnar committed 326 days ago

79231ecd

Check that exceptions are thrown when invalid types are passed to `promote_types`.

chr1sj0nes committed 326 days ago

bf57b7af

Use `jax.jit` instead of `pjit`. This fixes an error after pjit deprecation.

yashk2810 committed 326 days ago

9bcfac65

[hijax] make hijax types work with stages.Compiled / .compile()

mattjj committed 326 days ago

a63a91c1

[Mosaic TPU] Add stochastic rounding support on TPU 7x.

megrez-yliu committed 326 days ago

cc856e6f

Set the default value of `jax_ragged_dot_use_ragged_dot_instruction` to True.

pravnar committed 326 days ago

18f96c90

Merge pull request #32539 from Qazalbash:fix-betainc-doc

Google-ML-Automation committed 326 days ago

4feb2984

Migrate custom_partitioning_test to jax.jit.

pschuh committed 326 days ago

3e0108ed

Deprecate `jax.experimental.pjit.pjit`. Please use `jax.jit` instead.

yashk2810 committed 326 days ago

d339b6e8

Fix lax_test.RaggedTest

pravnar committed 326 days ago

d2142354

[buildcop] Disable tsan for `array_interoperability_test`

jakeharmon8 committed 326 days ago

13d55bd5

[buildcop] Disable asan for `array_interoperability_test` on GPU

jakeharmon8 committed 326 days ago

e53b164d

Merge pull request #32583 from mattjj:vjp3-accum-bug-fixes

Google-ML-Automation committed 326 days ago

51701b06

Disable asan for `mgpu_torch_test`

jakeharmon8 committed 327 days ago

09dfa032

[mutable-arrays] fix bugs with vjp3 exposed by #32470

mattjj committed 327 days ago

5ff9d7b1

Use abstract_mesh instead of concrete_mesh when binding mesh to SDS (when pspec is passed to SDS constructor). Also refactor device_assignment extraction logic in pxla.py.

yashk2810 committed 327 days ago

44a016be

Newer Older