ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

8ff5d334

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

Google-ML-Automation committed 1 year ago

2b510c0f

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

Google-ML-Automation committed 1 year ago

c31ad6dd

Remove version checks for ml_dtypes >= 0.5.0

apivovarov committed 1 year ago

cdb399a1

[Pallas] Support all integer casting cases.

justinjfu committed 1 year 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 1 year ago

ddc26812

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

yashk2810 committed 1 year ago

9b07a164

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

justinjfu committed 1 year ago

73ada4b9

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

Google-ML-Automation committed 1 year ago

ad0cebc4

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

Google-ML-Automation committed 1 year ago

917f902d

[doc] add docs for jax.lax.reduce_window

jakevdp committed 1 year ago

3b86ff23

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

Google-ML-Automation committed 1 year ago

b2a494f4

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

mattjj committed 1 year ago

c9f6bcdd

Merge pull request #30224 from jakevdp:openxla-urls

Google-ML-Automation committed 1 year ago

26200c50

Mark the jax.util submodule as deprecated

jakevdp committed 1 year ago

2d74566e

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

jakevdp committed 1 year ago

e86a2079

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

superbobry committed 1 year ago

e0443013

Merge pull request #30179 from carlosgmartin:lax_reduce_window_defaults

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

6f9240bd

Bump `jax_export_calling_convention_version` to 10

tomnatan30 committed 1 year ago

af55a484

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

jburnim committed 1 year ago

7972b928

Raise exception instead of returning NotImplementedError.

ZixuanJiang committed 1 year ago

be91cc1c

Fix to allow slicing into replicated dims

oliverdutton-iso committed 1 year ago

c90cfbe9

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

tomnatan30 committed 1 year ago

e0fe3bcb

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

bchetioui committed 1 year ago

26ea5d61

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

Google-ML-Automation committed 1 year ago

5eaef1ba

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

bchetioui committed 1 year ago

bd18922d

Merge pull request #30211 from gnecula:poly_triu_indices

Google-ML-Automation committed 1 year ago

ce48f711

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

bchetioui committed 1 year ago

83d89b7a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

f48cef57

Older