ROCm/jax

Pull Requests Commits

[ROCm] Remove redundant --config=rocm_base from wheel build command

mminutoli committed 35 days ago

a17d7209

Improve the lowering and printing of multi-ref transforms.

IvyZX committed 35 days ago

4ed7a0e7

[Pallas/Mosaic GPU] Make `reduce_scatter_mgpu` use warpgroup lowering semantics.

bchetioui committed 35 days ago

f05fc74b

[Pallas/Mosaic GPU] Make `collective_matmul_mgpu` use warpgroup lowering semantics.

bchetioui committed 35 days ago

55830196

[Pallas/Mosaic GPU] Make `all_gather_mgpu` use warpgroup lowering semantics.

bchetioui committed 35 days ago

44bdc919

[Pallas/Mosaic GPU] Make `hopper_mixed_type_matmul_mgpu` use warpgroup lowering semantics.

bchetioui committed 35 days ago

819def46

Merge pull request #37805 from samanklesaria:chisquare_checks

Google-ML-Automation committed 35 days ago

72b57290

Update out_sharding in _maxwell to match the new shape and explicitly mark the added dimension as replicated

yashk2810 committed 35 days ago

e556ec31

Merge pull request #37361 from vfdev-5:jrand-api--maxwell-out_sharding

Google-ML-Automation committed 35 days ago

b33092ef

Merge pull request #37797 from samanklesaria:uncomment_dtype_tests

Google-ML-Automation committed 35 days ago

9ea706ee

Speed up DynamicJaxprTracer.__init__.

hawkinsp committed 35 days ago

42dd8f75

Make JAX control deps compile on CPU.

mwhittaker committed 35 days ago

c6c8524f

Add threefry4x32 PRNG impl

Jake VanderPlas committed 35 days ago

5923cd99

Add check_all_safe_to_cast to chisquare

samanklesaria committed 35 days ago

e2940b20

Merge pull request #37740 from ROCm:switch_to_hermetic_llvm_upstream

Google-ML-Automation committed 35 days ago

3b8c84f0

Merge pull request #37416 from vfdev-5:jrand-api--generalized_normal-out_sharding

Google-ML-Automation committed 35 days ago

2451c1d6

Add select_ref tests using Pallas async_remote_copy.

IvyZX committed 35 days ago

fa3db6a4

Uncomment dtype tests in random_lax_test

samanklesaria committed 35 days ago

8f9133ea

Add philox4x32 PRNG impl

Jake VanderPlas committed 35 days ago

0fa13e32

[Pallas/Mosaic GPU] Make `blackwell_matmul_mgpu` use warpgroup lowering semantics.

bchetioui committed 35 days ago

31967d78

[Pallas/Mosaic GPU] Make `blackwell_ragged_dot_mgpu` use warpgroup lowering semantics.

bchetioui committed 35 days ago

90455038

Add link to advanced profiler options in JAX profiling guide

Google-ML-Automation committed 35 days ago

b70ce88e

Merge pull request #37679 from Intel-tensorflow:ashahba/docker-env-improvements

Google-ML-Automation committed 35 days ago

410dbaee

[Mosaic GPU] Allow `is_known_divisible` to handle `arith.TruncIOp`.

bchetioui committed 36 days ago

e958534c

Update the compilation test to handle different order in the backend configs.

pifon2a committed 36 days ago

1b4799ff

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

Google-ML-Automation committed 36 days ago

571a0385

Added out_sharding to jax.random.generalized_normal sampler

vfdev-5 committed 36 days ago

c626a86f

Added out_sharding to jax.random.maxwell sampler

vfdev-5 committed 36 days ago

5b47e2b5

Switch to hermetic llvm for rocm

alekstheod committed 36 days ago

2988f3ef

Add collective tests for compute_on2

Google-ML-Automation committed 36 days ago

b53699b6

Older