Commits ROCm/jax - SemanticDiff

intro Mosaic to ROCm, bulk commit. Needs fixes

Arech8 committed 23 days ago

e98a4266

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

Google-ML-Automation committed 23 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 23 days ago

480b5481

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

Google-ML-Automation committed 23 days ago

0b8b8cea

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

ZixuanJiang committed 23 days ago

875a7d6d

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

Google-ML-Automation committed 23 days ago

dd96fee0

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

Google-ML-Automation committed 23 days ago

d3753d4b

[mutable-arrays] add empty_ref primitive

mattjj committed 23 days ago

9dad8d5a

[Mosaic:TPU] Improve tpu.bitcast_vreg folding

tlongeri committed 23 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 24 days ago

f037af0c

Allow unreduced/reduced to work when check_vma=False.

yashk2810 committed 24 days ago

11fe9d20

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

bythew3i committed 24 days ago

56145d0a

Merge pull request #34231 from jakevdp:index-validation

Google-ML-Automation committed 24 days ago

f774eda1

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

ybaturina committed 24 days ago

cbf8f019

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

yashk2810 committed 24 days ago

aaec8ae5

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

yashk2810 committed 24 days ago

5054b519

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

apaszke committed 24 days ago

a0b906be

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

apaszke committed 24 days ago

c7a3bb5a

Manually bump the used XLA version

apaszke committed 24 days ago

32ea9761

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

Google-ML-Automation committed 24 days ago

75866a73

[indexing] internal refactor of indexing strategies

jakevdp committed 25 days ago

b088d3a0

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

Google-ML-Automation committed 25 days ago

0e19c821

[Pallas/interpreter] Support interpreting non-trivial GPU grids.

Google-ML-Automation committed 25 days ago

d799c9cc

Expand and update XProf docs

mtsokol committed 25 days ago

335bd052

[Pallas:MGPU] Increase shard_count for the interpret test

apaszke committed 25 days ago

3ebe54ed

Make `sub_byte_element_size_in_bits` public

Google-ML-Automation committed 25 days ago

2acc272a

Merge pull request #31720 from ROCm:ann-test-fix

Google-ML-Automation committed 25 days ago

81d95401

[Pallas:SC] Skip a debug_check test failing on TPU7x

apaszke committed 25 days ago

e5aad81b

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

Google-ML-Automation committed 25 days ago

b1272723

[pmap] Speed up the new `jax.pmap` implementation.

danielsuo committed 25 days ago

cf763131