ROCm/jax

Pull Requests Commits

Fix test behavior.

yueshengys committed 112 days ago

54add9ec

Removed a few unused functions

superbobry committed 112 days ago

b58b083e

Merge pull request #34710 from jakevdp:ref-indexing

Google-ML-Automation committed 112 days ago

ab376494

Reverts f3ac8813699fba956b85702813795e36867b0833

Google-ML-Automation committed 112 days ago

fc1b3213

[refs] fix ref indexing corner case

jakevdp committed 112 days ago

c2899d0c

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

superbobry committed 112 days ago

f3ac8813

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

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

480b5481

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

Google-ML-Automation committed 113 days ago

0b8b8cea

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

ZixuanJiang committed 113 days ago

875a7d6d

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

Google-ML-Automation committed 113 days ago

dd96fee0

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

Google-ML-Automation committed 113 days ago

d3753d4b

[mutable-arrays] add empty_ref primitive

mattjj committed 113 days ago

9dad8d5a

[Mosaic:TPU] Improve tpu.bitcast_vreg folding

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

f037af0c

Allow unreduced/reduced to work when check_vma=False.

yashk2810 committed 113 days ago

11fe9d20

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

bythew3i committed 113 days ago

56145d0a

Merge pull request #34231 from jakevdp:index-validation

Google-ML-Automation committed 113 days ago

f774eda1

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

ybaturina committed 113 days ago

cbf8f019

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

yashk2810 committed 113 days ago

aaec8ae5

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

yashk2810 committed 113 days ago

5054b519

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

apaszke committed 113 days ago

a0b906be

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

apaszke committed 113 days ago

c7a3bb5a

Manually bump the used XLA version

apaszke committed 113 days ago

32ea9761

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

Google-ML-Automation committed 113 days ago

75866a73

[indexing] internal refactor of indexing strategies

jakevdp committed 114 days ago

b088d3a0

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

Google-ML-Automation committed 114 days ago

0e19c821

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

Google-ML-Automation committed 114 days ago

d799c9cc

Expand and update XProf docs

mtsokol committed 114 days ago

335bd052

[Pallas:MGPU] Increase shard_count for the interpret test

apaszke committed 114 days ago

3ebe54ed

Newer Older