ROCm/jax

Pull Requests Commits

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

chr1sj0nes committed 1 year ago

5b0a767d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

540541a3

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

Google-ML-Automation committed 1 year ago

4ca97ad7

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

mattjj committed 1 year ago

a092df90

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

34aa5e69

Reverts ad21b62bfec5560d4c612ed3c8412eb2d240468b

Google-ML-Automation committed 1 year ago

d4745b9b

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

chr1sj0nes committed 1 year ago

74977938

[mosaic_gpu] Add `cupti_no_finalize` profiler mode.

chr1sj0nes committed 1 year ago

fd0ac022

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

chr1sj0nes committed 1 year ago

396e3890

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

hawkinsp committed 1 year ago

55e40847

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

pravnar committed 1 year ago

2692c5ff

Merge pull request #27326 from hawkinsp:tsan

Google-ML-Automation committed 1 year ago

5ce49bd2

Increase the test timeout for tsan builds.

hawkinsp committed 1 year ago

93f3e4aa

Remove symlinking of xla_client.py.

hawkinsp committed 1 year ago

ba5be78c

Remove some long-stale version guards.

hawkinsp committed 1 year ago

e71bcde5

Ensure traceback correctness in error checking in AOT mode

ayaka14732 committed 1 year ago

520b44fc

Reverts be5713309521d5cf0d2252b9c8f1d38ab50952d1

bmzhao committed 1 year ago

53e8eac7

Allow forcing pallas forward compatibility for some backends

krishnaharidasan committed 1 year ago

e23069b3

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

Google-ML-Automation committed 1 year ago

4fdce200

Add support for configurable error checking categories

ayaka14732 committed 1 year ago

7dd78d97

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

Google-ML-Automation committed 1 year ago

a9aa2a95

Merge pull request #26624 from wenscarl:bump_ml_dtypes

Google-ML-Automation committed 1 year 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 1 year ago

3163fbaa

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

Google-ML-Automation committed 1 year ago

e4ddbd1e

Add AOT support for error checking

ayaka14732 committed 1 year ago

3bf2eea2

Add `ShardingTypeError` to all sharding rules in JAX

yashk2810 committed 1 year ago

40ce44d1

Reorder C++ imports (nanobind).

hawkinsp committed 1 year ago

02719548

Deprecated `jax.tree_util.build_tree`

superbobry committed 1 year ago

92f5d9ca

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

bchetioui committed 1 year ago

27b30190

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

bchetioui committed 1 year ago

59d25f46

Older