Commits ROCm/jax - SemanticDiff

[ROCm] Enable previously skipped pallas tests on gfx950

amd-jianli12 committed 27 days ago

bc2ec359

[emit_pipeline] Fix async prefetch for the no outputs case.

rdyro committed 27 days ago

533182d0

Merge pull request #40338 from teddytennant:partition-order-invert

Google-ML-Automation committed 27 days ago

438c014b

Fix partition, argpartition and top_k on signed integer minimums

teddytennant committed 27 days ago

5a51ddd7

[Pallas] Skip test_async_sc_tc_prefetch_vmem with libtpu < 0.0.49.

hawkinsp committed 28 days ago

5df3e770

[Pallas][SC] Remove needs_layout_passes=False from debug_print tests

Google-ML-Automation committed 28 days ago

fc7fe71f

Lower shard_map in partial explicit manual mode correctly by using the aval.sharding.spec instead of `in_specs` since the former contain more info (explicit mesh axes too) while in_specs only refer to axes passed in axis_names.

yashk2810 committed 28 days ago

ae0ca5b8

Acquire a global lock in `jax.experimental.colocated_python.serialization._deserialize` to synchronize `cloudpickle.loads` calls across multiple threads.

Google-ML-Automation committed 28 days ago

20c6952b

Fix a CI test failure on Windows.

hawkinsp committed 28 days ago

e2e40d41

[Pallas][SC] Lower plsc.bitcast to tpu.bitcast when layout passes are enabled

Google-ML-Automation committed 28 days ago

cf8a8651

Merge pull request #40771 from jakevdp:pyrefly-update

Google-ML-Automation committed 28 days ago

7f3ebae3

[ci] update pyrefly to v1.2.0

jakevdp committed 28 days ago

013faf7d

Add strict_in_out (bool or tuple[bool, bool] representing strict_in, strict_out) to program_order to force first eqn's input and last eqn's output to be opt_barrier'd so that the first eqn and last eqn can't be sunk outside program_order(True) scope

yashk2810 committed 28 days ago

444168d6

Integrate LLVM at llvm/llvm-project@0038be81b8b5

Google-ML-Automation committed 28 days ago

53bb8f2a

Remove needs_layout_passes=False from tpu_sparsecore_pallas_distributed_test

Google-ML-Automation committed 28 days ago

43940433

Add a batched version of jax.copy_to_host_async.

Google-ML-Automation committed 28 days ago

bf979cc0

[Pallas/Mosaic GPU:Torch] Remove `notap` tag from Torch-related tests and migrate `mgpu_torch_test.py` to `plgpu.kernel`

bchetioui committed 28 days ago

f51644b9

[Pallas/Mosaic GPU] Migrate more downgraded GMEM stores to `plgpu.store`.

bchetioui committed 28 days ago

f89d4a9b

[Mosaic TPU] Fix some convolution failures.

yueshengys committed 28 days ago

2196def6

Merge pull request #40179 from zhenying-liu:cupti-13.3-v2

Google-ML-Automation committed 28 days ago

8bcf54a6

[fuser] decouple hardware dependencies and move TPU lowering rule to mosaic

rdyro committed 28 days ago

9ae0101f

Merge pull request #40572 from aybchan:hlo-collective-compat

Google-ML-Automation committed 28 days ago

8f29cb71

Add a batched version of jax.copy_to_host_async.

Google-ML-Automation committed 28 days ago

f4cf58f8

Document NaN ordering for descending sorts in jax.numpy.sort and argsort.

mkuperst committed 28 days ago

4f3d2c79

[Pallas/Mosaic GPU] Add a `plgpu.store` API.

bchetioui committed 28 days ago

ce2e28c3

Add tests for sort and reshape offloading in JAX memories_test.

Google-ML-Automation committed 28 days ago

5ce7a3f2

Optimize large JAX Mesh construction

Google-ML-Automation committed 28 days ago

0e319e53

internal change

Google-ML-Automation committed 28 days ago

3554f2d9

[Mosaic:GPU] Add support for multi-process execution with AbstractMesh in Mosaic GPU.

PatriosTheGreat committed 28 days ago

c18c1ad7

Reverts e592a247ba6058037f0d6c539df84b585d9d1338

thcmbs committed 28 days ago

45dbb343