ROCm/jax

Pull Requests Commits

Reenable tsan suppression, mark some tests as thread-unsafe.

hawkinsp committed 285 days ago

381f1109

Update XLA dependency to use revision

Google-ML-Automation committed 285 days ago

01904363

[Mosaic GPU] Add support for f16 Blackwell MMA accumulation

apaszke committed 285 days ago

a3e6c6ef

[mosaic_gpu] Add `Cupti` profiler class.

chr1sj0nes committed 285 days ago

0c38368b

Move PGLE documentation to JAX docs.

Google-ML-Automation committed 285 days ago

4da1faf5

[pallas] Add support for `split` (into two equal parts) in Triton lowering.

chr1sj0nes committed 285 days ago

a2475a66

[jax] Add `ndim` and `size` properties to `TransformedRef`.

chr1sj0nes committed 285 days ago

5b0a767d

Update XLA dependency to use revision

Google-ML-Automation committed 286 days ago

540541a3

Merge pull request #27353 from mattjj:remat-while-loop-fix

Google-ML-Automation committed 286 days ago

4ca97ad7

fix a linearize-of-remat-of-while_loop-fixpoint bug

mattjj committed 286 days ago

a092df90

Update XLA dependency to use revision

Google-ML-Automation committed 287 days ago

34aa5e69

Reverts ad21b62bfec5560d4c612ed3c8412eb2d240468b

Google-ML-Automation committed 287 days ago

d4745b9b

[pallas] Add support for `DotAlgorithmPreset.BF16_BF16_F32_X{6,9}` in Triton lowering.

chr1sj0nes committed 287 days ago

74977938

[mosaic_gpu] Add `cupti_no_finalize` profiler mode.

chr1sj0nes committed 287 days ago

fd0ac022

[pallas] Add `_zeros[_like]` and `_ones[_like]` utility functions in Triton lowering.

chr1sj0nes committed 287 days ago

396e3890

[JAX] [XLA:Python] Migrate xla_extension and its type stubs into jaxlib.

hawkinsp committed 287 days ago

55e40847

Lower lax.ragged_dot_general to chlo.ragged_dot in some cases on tpu.

pravnar committed 287 days ago

2692c5ff

Merge pull request #27326 from hawkinsp:tsan

Google-ML-Automation committed 287 days ago

5ce49bd2

Increase the test timeout for tsan builds.

hawkinsp committed 287 days ago

93f3e4aa

Remove symlinking of xla_client.py.

hawkinsp committed 287 days ago

ba5be78c

Remove some long-stale version guards.

hawkinsp committed 287 days ago

e71bcde5

Ensure traceback correctness in error checking in AOT mode

ayaka14732 committed 287 days ago

520b44fc

Reverts be5713309521d5cf0d2252b9c8f1d38ab50952d1

bmzhao committed 287 days ago

53e8eac7

Allow forcing pallas forward compatibility for some backends

krishnaharidasan committed 287 days ago

e23069b3

Add logit soft-capping support to the ragged paged attention Pallas kernel.

Google-ML-Automation committed 287 days ago

4fdce200

Add support for configurable error checking categories

ayaka14732 committed 287 days ago

7dd78d97

[Pallas] Fixes scalar prefetch in TPU interpret mode.

jburnim committed 287 days ago

37b5066d

Merge pull request #27308 from vfdev-5:tsan-ci-larger-timeout

Google-ML-Automation committed 287 days ago

a9aa2a95

Merge pull request #26624 from wenscarl:bump_ml_dtypes

Google-ML-Automation committed 287 days ago

7c53a9d2

Add varying manual axes rules to `mul_p` and `convert_element_type_p`. There are 2 things that need to be added:

yashk2810 committed 287 days ago

3163fbaa

Newer Older