ROCm/jax

Pull Requests Commits

Experimental pallas-triton lowering of ragged_dot_general on GPU.

rdyro committed 187 days ago

9a9fc023

Merge pull request #36537 from jakevdp:version-support

Google-ML-Automation committed 187 days ago

4d108a60

Update NumPy/SciPy version support

jakevdp committed 187 days ago

08aedc7c

[maint] add named deprecation for array dtype conversion

jakevdp committed 187 days ago

e2f64152

[JAX] Support non-PjRt-compatible Arrays in JAX heap profile API

hyeontaek committed 187 days ago

8c197055

Remove the jax_to_tf build rule and the tfcompile folder

gnecula committed 187 days ago

33dae7ee

[JAX] Use IFRT `Array::ByteSize()` to compute the on-device

hyeontaek committed 187 days ago

b0befdc7

[Pallas TPU] Add support for KEY memory space in emit_pipeline

sharadmv committed 187 days ago

7a3536eb

[consts] Fix AOT in presence of simplified handling of closed over consts

gnecula committed 187 days ago

54decb4e

Enable JAX async collectives test on CPU.

mwhittaker committed 187 days ago

ac4f6423

Make `mat.empty` a property instead of a method

yashk2810 committed 187 days ago

ed417730

Convert ManualAxisType to a dataclass with frozen=True so that we can remove the custom `__setattr__` which is slow.

yashk2810 committed 187 days ago

a48f5158

[pallas:mosaic] Handle missing `accuracy=` in `erf` and a few other unary ops

superbobry committed 187 days ago

7b07f9ae

[consts] Skip test that is irrelevant under JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS

gnecula committed 187 days ago

3bedc7c9

Prevent using a real backend for compile-only if platform versions differ.

marksandler2 committed 187 days ago

75b49de7

Automated Code Change

junwhanahn committed 187 days ago

ed6eb479

`pretty_printer.Doc` now uses `use_cpp_class`

superbobry committed 187 days ago

29261d2d

Fix single operand einsums when the output sharding is unreduced.

yashk2810 committed 187 days ago

3dc85d87

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

Google-ML-Automation committed 187 days ago

3607942d

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

Google-ML-Automation committed 188 days ago

613fc9c8

Removed a few unused internal functions

superbobry committed 189 days ago

647eaac7

Update XLA dependency to use revision http://github.com/openxla/xla/commit/5797c796dc2bcf1ace4b7852c8e0e393c0812317

Google-ML-Automation committed 189 days ago

54a9d3bb

Merge pull request #36506 from jax-ml:mattjj-patch-1

Google-ML-Automation committed 190 days ago

3ca9aeb2

update citation.bib

mattjj committed 190 days ago

Verified 626f3589

update author list

mattjj committed 190 days ago

Verified 3109b424

Merge pull request #36502 from mattjj:get-transpose-fix

Google-ML-Automation committed 190 days ago

4435b643

Merge pull request #36505 from jax-ml:sh_doc

Google-ML-Automation committed 190 days ago

723a303a

Fix sharding doc

yashk2810 committed 190 days ago

c8790940

[mutable-arrays] fix get transpose bug

mattjj committed 190 days ago

27f1ce6a

[jaxlib] Added all missing .pyis for jaxlib native extensions

superbobry committed 190 days ago

52669c9c

Newer Older