ROCm/jax

Pull Requests Commits

intro Mosaic to ROCm, bulk commit. Needs fixes

Arech8 committed 2 days ago

e98a4266

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

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

480b5481

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

Google-ML-Automation committed 2 days ago

0b8b8cea

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

ZixuanJiang committed 2 days ago

875a7d6d

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

Google-ML-Automation committed 3 days ago

dd96fee0

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

Google-ML-Automation committed 3 days ago

d3753d4b

[mutable-arrays] add empty_ref primitive

mattjj committed 3 days ago

9dad8d5a

[Mosaic:TPU] Improve tpu.bitcast_vreg folding

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

f037af0c

Allow unreduced/reduced to work when check_vma=False.

yashk2810 committed 3 days ago

11fe9d20

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

bythew3i committed 3 days ago

56145d0a

Merge pull request #34231 from jakevdp:index-validation

Google-ML-Automation committed 3 days ago

f774eda1

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

ybaturina committed 3 days ago

cbf8f019

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

yashk2810 committed 3 days ago

aaec8ae5

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

yashk2810 committed 3 days ago

5054b519

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

apaszke committed 3 days ago

a0b906be

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

apaszke committed 3 days ago

c7a3bb5a

Manually bump the used XLA version

apaszke committed 3 days ago

32ea9761

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

Google-ML-Automation committed 3 days ago

75866a73

[indexing] internal refactor of indexing strategies

jakevdp committed 4 days ago

b088d3a0

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

Google-ML-Automation committed 4 days ago

0e19c821

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

Google-ML-Automation committed 4 days ago

d799c9cc

Expand and update XProf docs

mtsokol committed 4 days ago

335bd052

[Pallas:MGPU] Increase shard_count for the interpret test

apaszke committed 4 days ago

3ebe54ed

Make `sub_byte_element_size_in_bits` public

Google-ML-Automation committed 4 days ago

2acc272a

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

Google-ML-Automation committed 4 days ago

81d95401

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

apaszke committed 4 days ago

e5aad81b

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

Google-ML-Automation committed 4 days ago

b1272723

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

danielsuo committed 4 days ago

cf763131

Older