ROCm/jax

Pull Requests Commits

Use blas version of complex numbers for expansion of templated GetrfBatched function

hsharsha committed 1 year ago

22414da4

Merge pull request #215 from ROCm/5.0_xla_sha_fix

zahiqbal committed 1 year ago

Verified b6843ff8

Update workspace.bzl

Ruturaj4 committed 1 year ago

Verified 1c9f9073

Update XLA hash for 0.5.0 branch (#212)

Ruturaj4 committed 1 year ago

Verified cdf1ef5a

Merge pull request #211 from ROCm/jax_5.0_xla_sha_fix

zahiqbal committed 1 year ago

Verified 81ac86ba

Update XLA hash for 0.5.0 branch

Ruturaj4 committed 1 year ago

b51b0c76

Release JAX 0.5.0

hawkinsp committed 1 year ago

c25fb92c

Reverts f1b894d14a28ac22a037fb79177b991275c75a18

Google-ML-Automation committed 1 year ago

a527aba6

[sharding_in_types] Error out for reshape for splits like this: `(4, 6, 8)` -> `(4, 4, 2, 6)`

yashk2810 committed 1 year ago

ce85b898

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7cac76d3

[Mosaic GPU] Delete unused declarations of `mosaic_gpu_memcpy_async_h2d`.

bchetioui committed 1 year ago

d3be190e

[mosaic_gpu] Added a serialization pass

superbobry committed 1 year ago

d34c40f6

[sharding_in_types] Rename `.at[...].get(out_spec)` to `.at[...].get(out_sharding)`.

yashk2810 committed 1 year ago

af667199

Rename out_type -> out_sharding parameter on einsum

yashk2810 committed 1 year ago

97cd7483

Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.

yashk2810 committed 1 year ago

49224d6c

[Mosaic TPU] Use large to compact 2nd minor retiling for conversions going both ways

apaszke committed 1 year ago

bd22bfef

Add quotes to pip commands in docs around option install for zsh

rdyro committed 1 year ago

aa9cea0a

Store GCS upload URI as a step output

nitins17 committed 1 year ago

5e52031d

Allow resharding between tokens on a single device

pschuh committed 1 year ago

f2f552c1

[sharding_in_types] If an indexing operation hits into `gather_p`, error out saying to use `.at[...].get(out_spec=...)` instead.

yashk2810 committed 1 year ago

b23c4237

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

994c3f59

[Mosaic] Fix infer/apply extensions.

WindQAQ committed 1 year ago

5c020ee3

Make `result_handler` of `_DeferredShardArg` a method instead of a property. Also play some code golf.

yashk2810 committed 1 year ago

0df4475a

[Mosaic TPU] Improve support for int16->int32 casts in TPUv4

apaszke committed 1 year ago

8954e71d

[MosaicGPU] Cleanup imports in dialect_lowering.py

dimitar-asenov committed 1 year ago

5e27efd0

[Mosaic GPU][NFC] Clean up import to align with stylistic guidance.

bchetioui committed 1 year ago

6746d633

[Mosaic GPU] Add layout inference for splat `arith.ConstantOp`s and `vector.SplatOp`s.

bchetioui committed 1 year ago

d3bf2433

[MosaicGPU] Remove the single_thread context from top-level dialect code.

dimitar-asenov committed 1 year ago

24884071

[Mosaic GPU][NFC] Simplify and clean up layout inference tests to use `FuncOp`s.

bchetioui committed 1 year ago

3366c927

[sharding_in_types] Expand reshape's sharding rule to add support for the following cases:

yashk2810 committed 1 year ago

c6b5ac5c

Older