ROCm/jax

Pull Requests Commits

[ROCm] Enable ToeplitzSymmetricConstruction and condition number unit tests for ROCm.

tsrw2048 committed 188 days ago

fcdd02ad

Fix vmap + shard_map + eager + shit combo and enable the commented out test

yashk2810 committed 188 days ago

8504c167

[pmap] Fix PRNGKeyArray type in multihost array conversion.

danielsuo committed 188 days ago

81ec7435

Remove an old version guard that was obsoleted by a jaxlib version bump.

hawkinsp committed 188 days ago

47d54e45

Remove USE_NEW_TPU_CALLBACK_LOWERING and bump libtpu version.

hawkinsp committed 188 days ago

4afee754

add input-to-output forwarding optimization to remat-of-jit

mattjj committed 188 days ago

8224ddd4

[pallas:mgpu] Fix `plgpu.kernel` vmap rule when `scratch_shapes` is passed as a `Mapping`.

chr1sj0nes committed 188 days ago

ccbdbc3e

Pin filelock to 3.20.3 to fix pytest jobs.

belitskiy committed 188 days ago

fc719f88

Make repr of PartitionSpec should print `P` instead of `PartitionSpec` cause it's too long.

yashk2810 committed 188 days ago

18474582

Fix some C++ linter warnings.

hawkinsp committed 188 days ago

9edaba1c

Merge pull request #34975 from jakevdp:pyrefly-fixes-2

Google-ML-Automation committed 188 days ago

25c0325d

Merge pull request #35039 from ezhulenev:type-id-hell-0

Google-ML-Automation committed 188 days ago

d2ce951e

Make oob indexing work with multiple devices if the output is sharded. Fixes https://github.com/jax-ml/jax/issues/35013

yashk2810 committed 189 days ago

4105150c

Merge pull request #34974 from ROCm:rocm-add-device-test

Google-ML-Automation committed 189 days ago

4916a429

[typing] more pyrefly fixes

jakevdp committed 189 days ago

29633fd9

Merge pull request #35036 from mattjj:allow-bf16-reductions-on-singleton-axes

Google-ML-Automation committed 189 days ago

66b3b7f5

[pallas] Use `compiler_params=` instead of `backend=` to disambiguate between Mosaic GPU and Triton

superbobry committed 189 days ago

daecb7e2

Assign explcit type ids to transpose plan caches

ezhulenev committed 189 days ago

d6fd55d1

Merge pull request #34574 from magaonka-amd:enable-variadic-reduce-window-rocm-upstream

Google-ML-Automation committed 189 days ago

ce5fb9e7

Merge pull request #34500 from magaonka-amd:fix-keyerror-rocm-upstream

Google-ML-Automation committed 189 days ago

fe4bbbc0

Merge pull request #35004 from jakevdp:more-pyrefly-fixes

Google-ML-Automation committed 189 days ago

d01881ad

Merge pull request #34675 from ROCm:skip-change-rocm

Google-ML-Automation committed 189 days ago

f0ed696c

allow bf16 reductions on singleton axes cf #34939

mattjj committed 189 days ago

4d91568c

[typing] fix pyrefly errors in jax/_src/numpy

jakevdp committed 189 days ago

28101604

[mpmd] Integrate PipelineSchedule API into jit lowering flow

Google-ML-Automation committed 189 days ago

5a3b1cc0

Merge pull request #34474 from ROCm:rocm_gesvdj_support

Google-ML-Automation committed 189 days ago

bad6b6ea

Merge pull request #35010 from jakevdp:pyrefly-fixes-3

Google-ML-Automation committed 189 days ago

38d3d87e

gulsumgudukbay committed 189 days ago

ebb3faba

address PR comment to remove map

gulsumgudukbay committed 189 days ago

0e83d5af

Update skip reasons for all tests to reflect they are <device> only instead of getting not supported on <device>

gulsumgudukbay committed 189 days ago

a069474f

Older