ROCm/jax

Pull Requests Commits

Add --no-cache option to the ci_build file

Ruturaj4 committed 352 days ago

64a2ab95

Update XLA hash for 0.5.0 branch (#212)

Ruturaj4 committed 356 days ago

Verified cdf1ef5a

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

zahiqbal committed 357 days ago

Verified 81ac86ba

Update XLA hash for 0.5.0 branch

Ruturaj4 committed 357 days ago

b51b0c76

Release JAX 0.5.0

hawkinsp committed 362 days ago

c25fb92c

Reverts f1b894d14a28ac22a037fb79177b991275c75a18

Google-ML-Automation committed 362 days ago

a527aba6

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

yashk2810 committed 362 days ago

ce85b898

Update XLA dependency to use revision

Google-ML-Automation committed 362 days ago

7cac76d3

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

bchetioui committed 363 days ago

d3be190e

[mosaic_gpu] Added a serialization pass

superbobry committed 363 days ago

d34c40f6

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

yashk2810 committed 363 days ago

af667199

Rename out_type -> out_sharding parameter on einsum

yashk2810 committed 363 days ago

97cd7483

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

yashk2810 committed 363 days ago

49224d6c

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

apaszke committed 363 days ago

bd22bfef

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

rdyro committed 363 days ago

aa9cea0a

Store GCS upload URI as a step output

nitins17 committed 363 days ago

5e52031d

Allow resharding between tokens on a single device

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

b23c4237

Update XLA dependency to use revision

Google-ML-Automation committed 363 days ago

994c3f59

[Mosaic] Fix infer/apply extensions.

WindQAQ committed 363 days ago

5c020ee3

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

yashk2810 committed 363 days ago

0df4475a

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

apaszke committed 363 days ago

8954e71d

[MosaicGPU] Cleanup imports in dialect_lowering.py

dimitar-asenov committed 363 days ago

5e27efd0

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

bchetioui committed 363 days ago

6746d633

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

bchetioui committed 363 days ago

d3bf2433

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

dimitar-asenov committed 363 days ago

24884071

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

bchetioui committed 363 days ago

3366c927

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

yashk2810 committed 363 days ago

c6b5ac5c

[MosaicGPU] Move `gpu_address_space_to_nvptx` inside `utils.py` and use it.

dimitar-asenov committed 363 days ago

ce03cf97

[Mosaic TPU] Add support for packing to 16-bit integers on TPUv4

apaszke committed 363 days ago

ef4dbd9c

Older