ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 208 days ago

8ff5d334

[Pallas] Physicalize fusion-dtype input avals before evaluating expressions.

Google-ML-Automation committed 208 days ago

2b510c0f

- Add more items in the ragged paged attention auto tuning table

Google-ML-Automation committed 208 days ago

c31ad6dd

Remove version checks for ml_dtypes >= 0.5.0

apivovarov committed 208 days ago

cdb399a1

[Pallas] Support all integer casting cases.

justinjfu committed 208 days ago

7507a3c0

[Mosaic GPU] Raise an error if TMEM allocations exceeds the limit. Currently, the kernel will hang which is difficult to debug.

justinjfu committed 208 days ago

ddc26812

Add a test mentioned in https://github.com/jax-ml/jax/issues/29883 to JAX test suite

yashk2810 committed 208 days ago

9b07a164

[Pallas][Mosaic GPU] Thread the python kernel function name to the MLIR kernel name.

justinjfu committed 208 days ago

73ada4b9

Merge pull request #30225 from jakevdp:lax-window-docs

Google-ML-Automation committed 208 days ago

ad0cebc4

Merge pull request #30227 from jakevdp:jax-util-dep

Google-ML-Automation committed 208 days ago

917f902d

[doc] add docs for jax.lax.reduce_window

jakevdp committed 208 days ago

3b86ff23

Merge pull request #30231 from mattjj:scan-linearize-instantiate

Google-ML-Automation committed 208 days ago

b2a494f4

[direct-linearize] add `instantiate` to linearize_jaxpr, fixpoint test

mattjj committed 208 days ago

c9f6bcdd

Merge pull request #30224 from jakevdp:openxla-urls

Google-ML-Automation committed 208 days ago

26200c50

Mark the jax.util submodule as deprecated

jakevdp committed 208 days ago

2d74566e

[doc] replace tensorflow.org/xla with openxla.org/xla

jakevdp committed 208 days ago

e86a2079

[pallas:mosaic] Log an error when an `AsyncCopyDescriptor` was unused

superbobry committed 208 days ago

e0443013

Merge pull request #30179 from carlosgmartin:lax_reduce_window_defaults

Google-ML-Automation committed 208 days ago

4a1e4799

Expand gather_sharding_rule to allow cases where operand dims in start_index_map are actually full_slices (and hence are not actually slicing into operand)

oliverdutton-iso committed 208 days ago

6f9240bd

Bump `jax_export_calling_convention_version` to 10

tomnatan30 committed 208 days ago

af55a484

[Pallas] Adds context manager for forcing TPU interpret mode.

jburnim committed 208 days ago

7972b928

Raise exception instead of returning NotImplementedError.

ZixuanJiang committed 208 days ago

be91cc1c

Fix to allow slicing into replicated dims

oliverdutton-iso committed 208 days ago

c90cfbe9

Add missing not-implemented sharding rule in `third_party/py/jax/_src/cudnn/fused_attention_stablehlo.py`.

tomnatan30 committed 208 days ago

e0fe3bcb

[Mosaic GPU] Remove unused `MosaicGPU_WGMMALayout` from the dialect.

bchetioui committed 209 days ago

26ea5d61

Merge pull request #30199 from andportnoy:mosaic-gpu-improve-runtime-error-messages

Google-ML-Automation committed 209 days ago

5eaef1ba

[Mosaic GPU] Add `Reduce` as a new constructor of `Expression`s.

bchetioui committed 209 days ago

bd18922d

Merge pull request #30211 from gnecula:poly_triu_indices

Google-ML-Automation committed 209 days ago

ce48f711

[Pallas/Mosaic GPU] Introduce the `print_layout` primitive.

bchetioui committed 209 days ago

83d89b7a

Update XLA dependency to use revision

Google-ML-Automation committed 209 days ago

f48cef57

Older