Commits ROCm/jax - SemanticDiff

[Mosaic] Add tpu.enqueue_indirect_dma verification test on TC

naummo committed 176 days ago

aebc56d3

Update XLA dependency to use revision

Google-ML-Automation committed 176 days ago

3b5bc35b

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

Jake VanderPlas committed 177 days ago

ee1037e1

Merge pull request #29943 from jakevdp:core-deps

Google-ML-Automation committed 177 days ago

43ecae00

Replace old heartbeat options with simplified heartbeat options.

mwhittaker committed 177 days ago

7d412a4e

Merge pull request #29930 from jakevdp:device-doc

Google-ML-Automation committed 177 days ago

9f9cc269

Improve error message when collective APIs are called without a shard_map

yashk2810 committed 177 days ago

ae792499

[dep] remove some finalized deprecations in jax.core

jakevdp committed 177 days 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 177 days 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 177 days ago

28fd6007

Document out_sharding in jax.numpy functions

jakevdp committed 177 days 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 177 days ago

c51c17c1

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

Jake VanderPlas committed 177 days ago

f7b5ae9e

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

Google-ML-Automation committed 177 days ago

e51b63ea

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

bchetioui committed 177 days 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 177 days 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 177 days ago

778a95b4

Update XLA dependency to use revision

Google-ML-Automation committed 178 days 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 178 days ago

3e8255fc

Add out_sharding to `jnp.arange`

yashk2810 committed 178 days ago

d241f6d6

Skip test when there aren't enough GPUs.

mwhittaker committed 178 days 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 178 days ago

52c0ed39

[JAX] Add a cache around HLO lowering rules.

hawkinsp committed 178 days ago

618460e8

Merge pull request #29433 from mwhittaker:disallow_vmap_aliasing

Google-ML-Automation committed 178 days ago

da3e738e

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

justinjfu committed 178 days ago

b6f53f0e

Fix JAX `scatter` docs.

zacmustin committed 178 days ago

d1dc79d1

Remove the implementation of infeed and outfeed.

hawkinsp committed 178 days ago

2fe18cbd

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

Jake VanderPlas committed 178 days ago

ab6182a8

Disallow aliased mutable array arguments to vmap.

mwhittaker committed 178 days ago

5f3e9502

Reverts 822c4f8534b7bd4ba9eea1d333fb9257af9285e7

mattjj committed 178 days ago

0f5c8cb9