ROCm/jax

Pull Requests Commits

[ROCm] Added ROCm tests to `device_test.py` unit test file.

tsrw2048 committed 2 days ago

2cf627b9

Merge pull request #34794 from jakevdp:pyrefly-fixes

Google-ML-Automation committed 2 days ago

7bbbc474

Update XLA dependency to use revision http://github.com/openxla/xla/commit/11991b81799e9f54a1eb6df417f089195f207c4f

Google-ML-Automation committed 2 days ago

e07d0645

[Mosaic GPU] Add missing assertion in `test_serialize_uses_bytecode_format`.

bchetioui committed 2 days ago

f1067f9a

Merge pull request #34939 from mattjj:disallow-f16-reductions

Google-ML-Automation committed 2 days ago

92506c17

[typing] some configuration & local fixes for pyrefly

jakevdp committed 2 days ago

db600bd9

add option to disallow f16 or bf16 reductions

mattjj committed 3 days ago

ba7321bf

Merge pull request #34949 from gnecula:shape_poly_test

Google-ML-Automation committed 3 days ago

b521b197

Merge pull request #34911 from superbobry:unused-variable

Google-ML-Automation committed 3 days ago

a008f1af

[pallas:sc] Extracted `lower_jaxpr_into_module`

superbobry committed 3 days ago

ac88a9f3

[RFC] Enabled unused-variable in Ruff's config

superbobry committed 3 days ago

5ebbdf2c

[Pallas/Mosaic GPU] Increase `shard_count` for `mosaic_gpu_test`.

bchetioui committed 3 days ago

9cccfe2c

[pallas] Fixed a type error in `_pallas_call_abstract_eval`

superbobry committed 3 days ago

30266a8b

[pallas:sc] Removed backward compatibility fix for empty grids

superbobry committed 3 days ago

be340a9a

Reverts 208d16833246ffc30f7f87cf3578272bf3b63890

superbobry committed 3 days ago

46945069

[pallas] `pallas_call` no longer supports multiple `compiler_params=`

superbobry committed 3 days ago

0bd8fca8

[shape_poly] Add a new test from a user question

gnecula committed 3 days ago

2df50916

Update XLA dependency to use revision http://github.com/openxla/xla/commit/9bc6ed9bc69b71c1d9b0dd9657bb6e78a66cbd1a

Google-ML-Automation committed 3 days ago

516e929c

Reverts bc7e997b88210f97029023885287f42de50c43f2

Google-ML-Automation committed 3 days ago

3d5fc32f

Reverts 13a22f6cd948a2ce43358aac3f86f079280b9b5f

Google-ML-Automation committed 3 days ago

208d1683

Add concat unreduced rule and split reduced rule

yashk2810 committed 3 days ago

2679eb33

[pallas] Deduped `get_interpreter_effects` between `core_map_p` and `pallas_call_p`

superbobry committed 3 days ago

13a22f6c

Merge pull request #34937 from jakevdp:argsort-dtype

Google-ML-Automation committed 3 days ago

5ac2bfdf

jnp.argsort: add optional dtype argument

jakevdp committed 3 days ago

2f8973ef

[pallas] `pallas_call` no longer supports multiple `compiler_params=`

superbobry committed 3 days ago

bc7e997b

Reverts e121d21893d46a613685363f1321dfe336ffbd2b

Google-ML-Automation committed 3 days ago

6654cd9c

Allow internal mutable array effect to be DCEd

sharadmv committed 3 days ago

54603e2e

[pmap] Update migrate pmap docs to include sharded reduction issues.

danielsuo committed 3 days ago

9bf9b863

[pmap] Cache jit wrappers in pmap's shard_map path to reduce dispatch overhead

danielsuo committed 3 days ago

280f6b4e

Pass unused recoverable option to coordination service.

mwhittaker committed 4 days ago

83e8cf58

Older