ROCm/jax

Pull Requests Commits

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

chr1sj0nes committed 299 days ago

5b0a767d

Update XLA dependency to use revision

Google-ML-Automation committed 299 days ago

540541a3

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

Google-ML-Automation committed 300 days ago

4ca97ad7

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

mattjj committed 300 days ago

a092df90

Update XLA dependency to use revision

Google-ML-Automation committed 300 days ago

34aa5e69

Reverts ad21b62bfec5560d4c612ed3c8412eb2d240468b

Google-ML-Automation committed 300 days ago

d4745b9b

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

chr1sj0nes committed 301 days ago

74977938

[mosaic_gpu] Add `cupti_no_finalize` profiler mode.

chr1sj0nes committed 301 days ago

fd0ac022

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

chr1sj0nes committed 301 days ago

396e3890

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

hawkinsp committed 301 days ago

55e40847

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

pravnar committed 301 days ago

2692c5ff

Merge pull request #27326 from hawkinsp:tsan

Google-ML-Automation committed 301 days ago

5ce49bd2

Increase the test timeout for tsan builds.

hawkinsp committed 301 days ago

93f3e4aa

Remove symlinking of xla_client.py.

hawkinsp committed 301 days ago

ba5be78c

Remove some long-stale version guards.

hawkinsp committed 301 days ago

e71bcde5

Ensure traceback correctness in error checking in AOT mode

ayaka14732 committed 301 days ago

520b44fc

Reverts be5713309521d5cf0d2252b9c8f1d38ab50952d1

bmzhao committed 301 days ago

53e8eac7

Allow forcing pallas forward compatibility for some backends

krishnaharidasan committed 301 days ago

e23069b3

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

Google-ML-Automation committed 301 days ago

4fdce200

Add support for configurable error checking categories

ayaka14732 committed 301 days ago

7dd78d97

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

Google-ML-Automation committed 301 days ago

a9aa2a95

Merge pull request #26624 from wenscarl:bump_ml_dtypes

Google-ML-Automation committed 301 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 301 days ago

3163fbaa

Merge pull request #27238 from MichaelHudgins:gpu-optional-presubmit

Google-ML-Automation committed 301 days ago

e4ddbd1e

Add AOT support for error checking

ayaka14732 committed 301 days ago

3bf2eea2

Add `ShardingTypeError` to all sharding rules in JAX

yashk2810 committed 301 days ago

40ce44d1

Reorder C++ imports (nanobind).

hawkinsp committed 301 days ago

02719548

Deprecated `jax.tree_util.build_tree`

superbobry committed 301 days ago

92f5d9ca

[Pallas/Mosaic GPU] Add lowering for WGMMA using warpgroup semantics.

bchetioui committed 301 days ago

27b30190

[Mosaic GPU] Add transform inference rule for `memref.load`.

bchetioui committed 301 days ago

59d25f46

Older