ROCm/jax

Pull Requests Commits

skip failing tests for ToT

gulsumgudukbay committed 153 days ago

07793745

Strip down workflow to execute only rocm piece.

mminutoli committed 153 days ago

368b60e1

Add ROCm pytest workflow to nightly/release CI

mminutoli committed 153 days ago

75420df1

Add pytest_rocm.yml workflow for ROCm testing

mminutoli committed 153 days ago

95223c6c

Add download-jax-rocm-wheels composite action

mminutoli committed 153 days ago

9778b54c

Add ROCm wheel patterns to install_wheels_locally.sh

mminutoli committed 153 days ago

72e857be

Complete run_pytest_rocm.sh for CI integration

mminutoli committed 153 days ago

5f869454

adding run_pytest_rocm.sh

gulsumgudukbay committed 153 days ago

8e541a61

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

yashk2810 committed 153 days ago

5054b519

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

apaszke committed 153 days ago

a0b906be

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

apaszke committed 153 days ago

c7a3bb5a

Manually bump the used XLA version

apaszke committed 153 days ago

32ea9761

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

Google-ML-Automation committed 153 days ago

75866a73

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

Google-ML-Automation committed 154 days ago

0e19c821

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

Google-ML-Automation committed 154 days ago

d799c9cc

Expand and update XProf docs

mtsokol committed 154 days ago

335bd052

[Pallas:MGPU] Increase shard_count for the interpret test

apaszke committed 154 days ago

3ebe54ed

Make `sub_byte_element_size_in_bits` public

Google-ML-Automation committed 154 days ago

2acc272a

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

Google-ML-Automation committed 154 days ago

81d95401

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

apaszke committed 154 days ago

e5aad81b

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

Google-ML-Automation committed 154 days ago

b1272723

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

danielsuo committed 154 days ago

cf763131

[pmap] Optimize host-local <> global conversions.

danielsuo committed 154 days ago

a2912b07

Fix: Use ASAN feature instead of config setting

yuriivcs committed 154 days ago

678acf28

[pmap] Push more into _cached_shard_map.

danielsuo committed 154 days ago

ee911642

[Mosaic] Expose `tpu::isGather`

naummo committed 154 days ago

2dbb22dd

Add env variables JAX_TRITON_PTXAS_DEVICE_DEBUG JAX_TRITON_PTXAS_EXTRA_FLAGS.

Google-ML-Automation committed 154 days ago

1f12c0a2

[mosaic:gpu] Use bytecode serialization.

hawkinsp committed 154 days ago

abae5e4e

[pmap] Remove any wrappings that have auxiliary values.

danielsuo committed 155 days ago

49561f67

[pmap] Cache computing effective mesh devices.

danielsuo committed 155 days ago

1cf66905

Older