Commits ROCm/jax - SemanticDiff

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 253 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 253 days ago

778a95b4

Update XLA dependency to use revision

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

3e8255fc

Add out_sharding to `jnp.arange`

yashk2810 committed 253 days ago

d241f6d6

Skip test when there aren't enough GPUs.

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

52c0ed39

[JAX] Add a cache around HLO lowering rules.

hawkinsp committed 253 days ago

618460e8

Merge pull request #29433 from mwhittaker:disallow_vmap_aliasing

Google-ML-Automation committed 253 days ago

da3e738e

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

justinjfu committed 253 days ago

b6f53f0e

Fix JAX `scatter` docs.

zacmustin committed 253 days ago

d1dc79d1

Remove the implementation of infeed and outfeed.

hawkinsp committed 253 days ago

2fe18cbd

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

Jake VanderPlas committed 253 days ago

ab6182a8

Disallow aliased mutable array arguments to vmap.

mwhittaker committed 254 days ago

5f3e9502

Reverts 822c4f8534b7bd4ba9eea1d333fb9257af9285e7

mattjj committed 254 days ago

0f5c8cb9

Make `axis_data.explicit_mesh_axis` dynamic whose value depends on the current mesh context.

yashk2810 committed 254 days ago

274117d2

Fix failing test because of stale StableHLO version.

mwhittaker committed 254 days ago

14ef0db6

Make sure the `carry` is varying across the same axes as `sorted_arr` when `scan` is called in `searchsorted`. Fixes https://github.com/jax-ml/jax/issues/29881

yashk2810 committed 254 days ago

2653c5b5

Merge pull request #29857 from AdamHillier:j-r-logistic

Google-ML-Automation committed 254 days ago

2496311d

Roll back https://github.com/jax-ml/jax/pull/29878 as the linter no longer warns in these cases.

Jake VanderPlas committed 254 days ago

7ac8181c

[Pallas TPU] Fix transformed ref's aval in pl.load/store and support multiple transforms in discharge rules.

bythew3i committed 254 days ago

7f3c9670

Merge pull request #29501 from Cjkkkk:jax_cudnn_9.10_sdpa_fp8

Google-ML-Automation committed 254 days ago

443b881a

Merge pull request #29898 from mattjj:no-more-strong-tracer-refs

Google-ML-Automation committed 254 days ago

822c4f85

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

Jake VanderPlas committed 254 days ago

fb833d77

simplify our jaxpr builder not to hold strong refs to tracers

mattjj committed 254 days ago

485b681b

[Mosaic GPU] Implement a skeleton for the new equation-system-driven layout inference.

bchetioui committed 254 days ago

6650663b

[Mosaic GPU] Factor parametrization of WGMMA tests

apaszke committed 254 days ago

dba6e224

[pallas:triton] Use `metadata` instead of `TritonCompilerParams.serialized_metadata`

superbobry committed 254 days ago

f754fb08

[Pallas][Mosaic GPU] Rename for_tensor_core to orders_tensor_core

apaszke committed 254 days ago

e05385d9

[Pallas:MGPU] Make TMEM reads and writes explicitly asynchronous

apaszke committed 254 days ago

967489d8