ROCm/jax

Pull Requests Commits

Skip version checking for ROCm plugins

Ruturaj4 committed 6 days ago

42d5a8f3

relax version requirements for ROCm Jax Plugin wheels, allow overriding using env var

JehandadKhan committed 6 days ago

82ed029f

[indexing] avoid printing full array repr in indexing-related error messages.

Jake VanderPlas committed 6 days ago

17b1bfe7

Fix test behavior.

yueshengys committed 6 days ago

54add9ec

Removed a few unused functions

superbobry committed 6 days ago

b58b083e

Merge pull request #34710 from jakevdp:ref-indexing

Google-ML-Automation committed 6 days ago

ab376494

Reverts f3ac8813699fba956b85702813795e36867b0833

Google-ML-Automation committed 6 days ago

fc1b3213

[refs] fix ref indexing corner case

jakevdp committed 6 days ago

c2899d0c

[mosaic] Avoid using the deprecated `output_memory_colors` in `CustomCallBackendConfig`

superbobry committed 6 days ago

f3ac8813

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

Google-ML-Automation committed 6 days ago

9489ca75

[Mosaic TPU] Support arbitrary offset during mask generation for 16-bit data in `SingleRowVRegBounds` by leveraging the 16-bit tpu.iota op.

yueshengys committed 6 days ago

480b5481

[Pallas/interpreter] Fix incorrect nesting of acquisitions of locks/condition variables.

Google-ML-Automation committed 6 days ago

0b8b8cea

Disable the `contracting2` test case from `test_dot_general`, following cl/859306670.

ZixuanJiang committed 7 days ago

875a7d6d

[Pallas/interpreter] Add an internal consistency check at the end of executing interpreted GPU kernels.

Google-ML-Automation committed 7 days ago

dd96fee0

Merge pull request #34713 from mattjj:empty-ref-p

Google-ML-Automation committed 7 days ago

d3753d4b

[mutable-arrays] add empty_ref primitive

mattjj committed 7 days ago

9dad8d5a

[Mosaic:TPU] Improve tpu.bitcast_vreg folding

tlongeri committed 7 days ago

88a07221

Move check_vma check inside pvary traceable instead of it's abstract_eval rule making it a complete no-op. Before pvary as a primitive was still in the jaxpr so it's transpose rule could be triggered.

yashk2810 committed 7 days ago

f037af0c

Allow unreduced/reduced to work when check_vma=False.

yashk2810 committed 7 days ago

11fe9d20

[Mosaic TPU] Fix tpu.bitcast_vreg fold rule when input def op is null.

bythew3i committed 7 days ago

56145d0a

Merge pull request #34231 from jakevdp:index-validation

Google-ML-Automation committed 7 days ago

f774eda1

Make `clone_main_xla` workflow parameter required to have an explicit control whether the main XLA should be cloned.

ybaturina committed 7 days ago

cbf8f019

Raise an error if src or dst passed to device_put is unreduced.

yashk2810 committed 7 days ago

aaec8ae5

Raise a NotImplementedError if unreduced/reduced is used with eager shard_map.

yashk2810 committed 7 days ago

5054b519

[Pallas:SC] Fix an incorrectly implemented test skip condition

apaszke committed 7 days ago

a0b906be

[Mosaic GPU] Fix a MLIR error when doing layout conversions involving f8 types

apaszke committed 7 days ago

c7a3bb5a

Manually bump the used XLA version

apaszke committed 7 days ago

32ea9761

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

Google-ML-Automation committed 7 days ago

75866a73

[indexing] internal refactor of indexing strategies

jakevdp committed 8 days ago

b088d3a0

Merge pull request #33881 from mtsokol:update-xprof-docs

Google-ML-Automation committed 8 days ago

0e19c821

Older