ROCm/jax

Pull Requests Commits

[ROCm] Add rocm support in nm.py lowering code

Ruturaj4 committed 2 years ago

d2ab42de

[Mosaic GPU] Make the profiler warpgroup aware

apaszke committed 2 years ago

9e1a4e8d

Moved op-specific tests from PallasCallTest to PallasOpsTest

superbobry committed 2 years ago

208fd2bb

[xla:cpu] Add support for outfeed to thunk runtime

ezhulenev committed 2 years ago

4cca2335

[jax:cpu] Add a test config to run JAX with new XLA:CPU runtime

ezhulenev committed 2 years ago

1aa2b6ee

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

95ed39bf

Merge pull request #21519 from andportnoy:aportnoy/ffi-build-capsule

a-googler committed 2 years ago

10e8a0b9

Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.

yashk2810 committed 2 years ago

12730280

[Mosaic] Document tpu.create_subelement_mask

tlongeri committed 2 years ago

8a1445a0

[Mosaic] Remove "support" for MAXNUMF vector reductions

tlongeri committed 2 years ago

e620acfa

Add spaces before dump instructions

Michael Levesque-Dion committed 2 years ago

2bfafe35

[pallas:Mosaic GPU] Capture stdout in mosaic tests for debugging.

cperivol committed 2 years ago

d99cd3e9

Add pycapsule utility to jax.extend.ffi

andportnoy committed 2 years ago

fa46bbf0

[Mosaic GPU] Move the matmul example runner away from the test harness

apaszke committed 2 years ago

a7e35c6b

`input_shardings` returned from the compiled executable should match the in_tree pre DCE. Otherwise calling it raises an error saying that the length of pre dce in_tree and post dce input_shardings don't match

yashk2810 committed 2 years ago

9d5ee61b

Merge pull request #21598 from gnecula:doctest

a-googler committed 2 years ago

23549cda

[Mosaic GPU] Support f16 wgmma accumulator type.

cperivol committed 2 years ago

726adc03

[JAX] Fail gracefully when an array with multiple shards is passed to make_array_from_single_device_arrays.

hawkinsp committed 2 years ago

b2d20891

Fixed a typo in `_pallas_call_impl`

superbobry committed 2 years ago

fe35acd4

Merge pull request #21593 from superbobry:pallas

a-googler committed 2 years ago

c59b08c8

[Mosaic GPU] Skip tests for Python versions < 3.10

apaszke committed 2 years ago

84dee133

Slightly rearranged NDIndexer.from_indices_shape and added missing tests

superbobry committed 2 years ago

befa10c1

[Pallas] Clean up MLIR compatibility check

tlongeri committed 2 years ago

7c471e25

[docs] Enable doctest for .md documentation files

gnecula committed 2 years ago

0223e830

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

a5e76130

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

e81c8260

Fixed the logging due to a recent change.

datenglin committed 2 years ago

20379c63

Copybara import of the project:

gnecula committed 2 years ago

be1e40dc

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

432159a9

[JAX] Use iota_reshape_dims and iota_transpose_perm in pxla, which is more efficient than tile_assignment_devices.

ZixuanJiang committed 2 years ago

dfd7d17c

Older