ROCm/jax

Pull Requests Commits

[ROCM] nits and fixes

Ruturaj4 committed 1 year ago

68265702

[JAX] change rocm build file for pjrt plugin

Ruturaj4 committed 1 year ago

bc0fa9e4

[ROCm] Bring-up pjrt support

Ruturaj4 committed 1 year ago

7d9aa651

[ROCm] Add hip specific checks in threefry test

Ruturaj4 committed 1 year ago

0ab9785e

[ROCm] skip failing profiling tests

Ruturaj4 committed 1 year ago

ad701f69

[ROCm]: Add support to continue on fail

Rahul Batra committed 1 year ago

5338c4c7

[Mosaic GPU] Make the profiler warpgroup aware

apaszke committed 1 year ago

9e1a4e8d

Moved op-specific tests from PallasCallTest to PallasOpsTest

superbobry committed 1 year ago

208fd2bb

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

ezhulenev committed 1 year ago

4cca2335

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

ezhulenev committed 1 year ago

1aa2b6ee

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

95ed39bf

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

a-googler committed 1 year 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 1 year ago

12730280

[Mosaic] Document tpu.create_subelement_mask

tlongeri committed 1 year ago

8a1445a0

[Mosaic] Remove "support" for MAXNUMF vector reductions

tlongeri committed 1 year ago

e620acfa

Add spaces before dump instructions

Michael Levesque-Dion committed 1 year ago

2bfafe35

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

cperivol committed 1 year ago

d99cd3e9

Add pycapsule utility to jax.extend.ffi

andportnoy committed 1 year ago

fa46bbf0

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

apaszke committed 1 year 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 1 year ago

9d5ee61b

Merge pull request #21598 from gnecula:doctest

a-googler committed 1 year ago

23549cda

[Mosaic GPU] Support f16 wgmma accumulator type.

cperivol committed 1 year ago

726adc03

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

hawkinsp committed 1 year ago

b2d20891

Fixed a typo in `_pallas_call_impl`

superbobry committed 1 year ago

fe35acd4

Merge pull request #21593 from superbobry:pallas

a-googler committed 1 year ago

c59b08c8

[Mosaic GPU] Skip tests for Python versions < 3.10

apaszke committed 1 year ago

84dee133

Slightly rearranged NDIndexer.from_indices_shape and added missing tests

superbobry committed 1 year ago

befa10c1

[Pallas] Clean up MLIR compatibility check

tlongeri committed 1 year ago

7c471e25

[docs] Enable doctest for .md documentation files

gnecula committed 1 year ago

0223e830

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a5e76130

Older