ROCm/jax

Pull Requests Commits

Skip tests on ROCm

Ruturaj4 committed 312 days ago

181b0944

Fix invalid lowerings for ROCm in Pallas (#223)

mrodden committed 314 days ago

Verified 6b351552

Use HIPBLAS_V2 (#222)

Ruturaj4 committed 323 days ago

Verified 1fc3d157

Merge pull request #215 from ROCm/5.0_xla_sha_fix

zahiqbal committed 331 days ago

Verified b6843ff8

Update workspace.bzl

Ruturaj4 committed 331 days ago

Verified 1c9f9073

Update XLA hash for 0.5.0 branch (#212)

Ruturaj4 committed 336 days ago

Verified cdf1ef5a

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

zahiqbal committed 337 days ago

Verified 81ac86ba

Update XLA hash for 0.5.0 branch

Ruturaj4 committed 337 days ago

b51b0c76

Release JAX 0.5.0

hawkinsp committed 342 days ago

c25fb92c

Reverts f1b894d14a28ac22a037fb79177b991275c75a18

Google-ML-Automation committed 342 days ago

a527aba6

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

yashk2810 committed 342 days ago

ce85b898

Update XLA dependency to use revision

Google-ML-Automation committed 342 days ago

7cac76d3

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

bchetioui committed 342 days ago

d3be190e

[mosaic_gpu] Added a serialization pass

superbobry committed 342 days ago

d34c40f6

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

yashk2810 committed 343 days ago

af667199

Rename out_type -> out_sharding parameter on einsum

yashk2810 committed 343 days ago

97cd7483

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

yashk2810 committed 343 days ago

49224d6c

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

apaszke committed 343 days ago

bd22bfef

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

rdyro committed 343 days ago

aa9cea0a

Store GCS upload URI as a step output

nitins17 committed 343 days ago

5e52031d

Allow resharding between tokens on a single device

pschuh committed 343 days 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 343 days ago

b23c4237

Update XLA dependency to use revision

Google-ML-Automation committed 343 days ago

994c3f59

[Mosaic] Fix infer/apply extensions.

WindQAQ committed 343 days ago

5c020ee3

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

yashk2810 committed 343 days ago

0df4475a

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

apaszke committed 343 days ago

8954e71d

[MosaicGPU] Cleanup imports in dialect_lowering.py

dimitar-asenov committed 343 days ago

5e27efd0

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

bchetioui committed 343 days ago

6746d633

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

bchetioui committed 343 days ago

d3bf2433

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

dimitar-asenov committed 343 days ago

24884071

Older