Commits ROCm/jax - SemanticDiff

[Mosaic] Add tpu.enqueue_indirect_dma verification test on TC

naummo committed 1 year ago

aebc56d3

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

3b5bc35b

Move `jax._src.image` to its own build rule

Jake VanderPlas committed 1 year ago

ee1037e1

Merge pull request #29943 from jakevdp:core-deps

Google-ML-Automation committed 1 year ago

43ecae00

Replace old heartbeat options with simplified heartbeat options.

mwhittaker committed 1 year ago

7d412a4e

Merge pull request #29930 from jakevdp:device-doc

Google-ML-Automation committed 1 year ago

9f9cc269

Improve error message when collective APIs are called without a shard_map

yashk2810 committed 1 year ago

ae792499

[dep] remove some finalized deprecations in jax.core

jakevdp committed 1 year ago

4c28eded

Allow GSPMDSharding constructor to take in device_list (xc.DeviceList) as input along with `Sequence[jax.Device]`. This prevents extremely slow `tuple(devices) -> DeviceList` conversion in the GSPMDSharding constructor.

yashk2810 committed 1 year ago

14b2c90f

If `jax_ragged_dot_use_ragged_dot_instruction == True` then use_default_lowering() should return False since default lowering is via a pattern match.

Google-ML-Automation committed 1 year ago

28fd6007

Document out_sharding in jax.numpy functions

jakevdp committed 1 year ago

acc158e0

Adds jax_ragged_dot_use_ragged_dot_instruction config/flag (only applicable to TPUs) that forces ragged_dot() lowering to use chlo.ragged_dot instruction.

Google-ML-Automation committed 1 year ago

c51c17c1

Move `jax._src.scipy` to its own build rule

Jake VanderPlas committed 1 year ago

f7b5ae9e

Merge pull request #29901 from harini-sridhar:patch-1

Google-ML-Automation committed 1 year ago

e51b63ea

[Mosaic GPU] Introduce `{Least,Most}ReplicatedExpression` in the equation system.

bchetioui committed 1 year ago

525a13fb

Preserve the sharding provided in `out_sharding` argument of `auto_axes`/or any other API when the computation has a single device in explicit mode

yashk2810 committed 1 year ago

5ab714bd

Fix a bug introduced in https://github.com/jax-ml/jax/pull/29910. Only take the `out_sharding` path in `_arange` if the sharding is not concrete.

yashk2810 committed 1 year ago

778a95b4

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

07c990da

Bring back `os.getenv` for `jax_dump_ir_to`, `jax_include_debug_info_in_dumps` and `jax_dump_ir_modes` because they are flags and not states. `env-var` by default only works for states.

yashk2810 committed 1 year ago

3e8255fc

Add out_sharding to `jnp.arange`

yashk2810 committed 1 year ago

d241f6d6

Skip test when there aren't enough GPUs.

mwhittaker committed 1 year ago

75b0d0a6

Allow decorator factory pattern for `smap` i.e. `@smap(in_axes=..., out_axes=..., axis_name=...)`. This is similar to `jax.shard_map`.

yashk2810 committed 1 year ago

52c0ed39

[JAX] Add a cache around HLO lowering rules.

hawkinsp committed 1 year ago

618460e8

Merge pull request #29433 from mwhittaker:disallow_vmap_aliasing

Google-ML-Automation committed 1 year ago

da3e738e

[Pallas TPU] Add support for per-input multiple buffering to emit_pipeline.

justinjfu committed 1 year ago

b6f53f0e

Fix JAX `scatter` docs.

zacmustin committed 1 year ago

d1dc79d1

Remove the implementation of infeed and outfeed.

hawkinsp committed 1 year ago

2fe18cbd

Move `jax._src.tpu` to its own build rule

Jake VanderPlas committed 1 year ago

ab6182a8

Disallow aliased mutable array arguments to vmap.

mwhittaker committed 1 year ago

5f3e9502

Reverts 822c4f8534b7bd4ba9eea1d333fb9257af9285e7

mattjj committed 1 year ago

0f5c8cb9