ROCm/jax

Pull Requests Commits

Rename the jax_shard_map_partitioning configuration flag to jax_partitioning.

yashk2810 committed 4 days ago

b82be556

Fix jnp.choose for narrow index dtypes and large choice counts

tanveer-arch committed 4 days ago

91a6f34f

[Mosaic GPU] Allow allocating up to 576 TMEM columns on sm_107f.

golechwierowicz committed 4 days ago

447f46e4

Align TMEM transfer layout inference with TMEMRef.load/store.

golechwierowicz committed 4 days ago

1ad1de75

Plumb the target architecture into layout inference derivation rules.

golechwierowicz committed 4 days ago

d430af30

[jax-buildcop] Stop worker subprocess cleanly before joining in `ProfilerTest.testSubprocessProfiling`

Google-ML-Automation committed 4 days ago

6dca3ace

[Mosaic GPU] Lift restriction on 128-lane A scales in sparse block-scaled MMA on sm_107f

golechwierowicz committed 4 days ago

f5981d98

Lift restriction on 128-lane A scales in collective M=256 block-scaled MMA on sm_107f.

golechwierowicz committed 4 days ago

f719a9a5

jax.scipy.fft.dct: use rfft for real input

zazabap committed 4 days ago

f64b1205

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

Google-ML-Automation committed 4 days ago

fd48df9b

[Mosaic:GPU] Parse collective attributes only once and store it in CustomCallResources during instantiation.

PatriosTheGreat committed 4 days ago

776c7a46

Reverts 1e384f92d5ce5a419a7c4ba4645cff88a7080c74

Google-ML-Automation committed 4 days ago

4473bda5

Remove `auto_axes` from `jnp.histogram`.

levskaya committed 4 days ago

ba831131

Preserve explicit batch shardings in `lax.linalg` `_tril`, `_triu`, and `_qr_lowering`.

levskaya committed 4 days ago

975a6059

Remove `auto_axes` from `jnp.repeat`.

levskaya committed 4 days ago

5b002e42

[jax-buildcop] Restore redirected file descriptor before reading captured output in `_capture_output`

Google-ML-Automation committed 4 days ago

130b5387

Add jax.scipy.special.kn for integer-order modified Bessel functions of the second kind.

Jake VanderPlas committed 4 days ago

4d47f474

[partitioning] Support sharded argmin and argmax.

rdyro committed 4 days ago

34af4049

[Tests] Skip `bfloat16` `testF8E4M3FNDotAlgorithmAccuracy` on CUDA 8.9.

belitskiy committed 4 days ago

afb7dc80

Document cross-compiling for Cloud TPU via the compilation cache.

Google-ML-Automation committed 4 days ago

9a274688

Add `out_sharding` to `jnp.take`.

levskaya committed 4 days ago

49a98685

Fix non-termination in `jnp.gcd` and `jnp.lcm` for minimum signed integers.

Jake VanderPlas committed 4 days ago

5a5bbe33

[partitioning] Support sharded lax.iota and random_clone passthrough.

rdyro committed 4 days ago

49aa8769

[jax-buildcop] Disable `mutable_array_test` on CPU under TSan

Google-ML-Automation committed 4 days ago

16ee1909

Add `mode="high"` to `jax.random.truncated_normal` to fix tail sampling accuracy.

Jake VanderPlas committed 4 days ago

19a93c08

Merge pull request #42419 from Intel-tensorflow:mraunak/C++20_oneapi

Google-ML-Automation committed 5 days ago

06390fc5

Add regression test for TPU gather grad bug.

Jake VanderPlas committed 5 days ago

2a415ea2

Integrate LLVM at llvm/llvm-project@df0a91f644ef

Google-ML-Automation committed 5 days ago

09e7414d

Upcast narrow signed integer indices before negative-index wrapping.

Jake VanderPlas committed 5 days ago

9d63d567

Avoid rank-changing slice/concat ops in `jax.random.gumbel` (`mode="highest"`).

henryhmko committed 5 days ago

1e384f92

Newer Older