ROCm/jax

Pull Requests Commits

[Mosaic] Rename dep name.

WindQAQ committed 312 days ago

5179642e

Fix convolution example (kernel should be OIHW, not IOHW).

Google-ML-Automation committed 312 days ago

1a57fdf7

[Mosaic GPU] Make the small WGMMA tile independent of transpose flags

apaszke committed 312 days ago

e9f95cc3

#sdy Add JAX backwards compatibility test.

bartchr808 committed 313 days ago

ed4a7bba

#sdy support JAX export tests when Shardy is enabled.

bartchr808 committed 313 days ago

ac493655

[mgpu] Forach in tiled layout.

cperivol committed 313 days ago

b9ebd918

[Mosaic GPU][NFC] Move the calculation of group strides into _validate_mma

apaszke committed 313 days ago

11e6cfbc

Merge pull request #26697 from gnecula:pp_aliased_var_names

Google-ML-Automation committed 313 days ago

bbadf990

[Mosaic GPU][NFC] Clean up the computation of group strides

apaszke committed 313 days ago

3038348f

Use the same name for aliased Vars when pretty-printing Jaxprs.

gnecula committed 313 days ago

a6c47d6f

Update XLA dependency to use revision

Google-ML-Automation committed 313 days ago

eee4d601

Add use_high_dynamic_range_gumbel flag which allows sampling gumbel such

pschuh committed 313 days ago

b8b690e5

[Mosaic GPU] Infer layouts (transforms) on memrefs that directly feed into the dialect wgmma op.

dimitar-asenov committed 314 days ago

3b305c66

[Mosaic GPU] Wire up the `slice_lengths` and `indices` operands in lowering of the MLIR dialect.

dimitar-asenov committed 314 days ago

c60ef5a2

Update XLA dependency to use revision

Google-ML-Automation committed 314 days ago

9d6bcd63

`PRNGKeyArray.aval` should have the correct logical sharding. This required refactoring code so that we don't hit recursion errors.

yashk2810 committed 314 days ago

53494ade

Update XLA dependency to use revision

Google-ML-Automation committed 315 days ago

e25caba9

Chnages for kernel export

Google-ML-Automation committed 315 days ago

2a1eeb0c

Reverts 10f6edeb496a2eec2a09c2c5cecbe4f8f02452ab

Anton Osokin committed 315 days ago

1f317663

Improved errors when indexing with floats

jakevdp committed 315 days ago

b2c45b8e

Add a profiler test for gpu run

Google-ML-Automation committed 315 days ago

48a55a6d

doc: fix description of logsumexp axis

jakevdp committed 315 days ago

c56e794a

Temporarily skip some more linalg sharding checks.

dfm committed 315 days ago

70024d22

Bump oldest supported libtpu to match the compatibility window (12 weeks)

kanglant committed 315 days ago

da7c90c4

Print the list of installed packages before running pytests

nitins17 committed 315 days ago

0ed42dcd

[sharding_in_types] `out_sharding` argument on einsum should only apply to the last einsum and not intermediate einsums.

yashk2810 committed 315 days ago

da1cc0a5

Merge pull request #26814 from dfm:remat-opt-custom-dce

Google-ML-Automation committed 315 days ago

10f6edeb

Merge pull request #26835 from skye:libtpu_version

Google-ML-Automation committed 315 days ago

6ee261e9

Add build targets for `jax-rocm-plugin` and `jax-rocm-pjrt` wheels.

Google-ML-Automation committed 315 days ago

8f57b816

Reimplement custom_vjp.optimize_remat using custom_dce.

dfm committed 316 days ago

bb9aed5e

Newer Older