ROCm/jax

Pull Requests Commits

add mosaic skips

Ruturaj4 committed 27 days ago

58efc147

Merge pull request #35332 from superbobry:pyrefly

Google-ML-Automation committed 27 days ago

b486aacf

Pin mpmath to v1.3.0.

danielsuo committed 27 days ago

83471837

[jaxlib] Bundle _{tpu,triton,mosaic_gpu}_ext.pyi with jaxlib

superbobry committed 27 days ago

22875809

[Mosaic TPU] Give more time for the changes to propagate to libtpu

apaszke committed 27 days ago

db34816c

Another batch of Pyrefly fixes

superbobry committed 27 days ago

e673505c

Merge pull request #35317 from superbobry:pyrefly

Google-ML-Automation committed 28 days ago

ad8c8c4f

Merge pull request #35323 from jakevdp:mpmath-fix

Google-ML-Automation committed 28 days ago

adc7836d

More Pyrefly fixes in Pallas

superbobry committed 28 days ago

0beb8372

[test] fix deprecated mpmath import

jakevdp committed 28 days ago

b5c92015

Merge pull request #33802 from jax-ml:dependabot/github_actions/actions/checkout-6.0.1

Google-ML-Automation committed 28 days ago

a58b181f

[pallas:mgpu] Fix handling of `AbstractRef` in `_commute_transform`.

chr1sj0nes committed 28 days ago

cd019717

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

Google-ML-Automation committed 28 days ago

f47f0c20

Merge pull request #35309 from superbobry:pyrefly

Google-ML-Automation committed 28 days ago

4d3395f0

[pyrefly] Fixed or suppressed type errors in Pallas backends

superbobry committed 28 days ago

9ce7295f

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

Google-ML-Automation committed 29 days ago

8dac162f

Merge pull request #35306 from mattjj:remat-jaxpr

Google-ML-Automation committed 29 days ago

dea31b6c

[hijax][remat3] add basic remat_jaxpr

mattjj committed 29 days ago

b17e41dd

Merge pull request #35304 from mattjj:remat-transform

Google-ML-Automation committed 29 days ago

53d6b350

[hijax][remat3] start implementing remat on hijax

mattjj committed 29 days ago

5f43a720

Merge pull request #35302 from mattjj:hijax-dce-fix

Google-ML-Automation committed 29 days ago

0cff9fe3

[hijax] fix followup on #35301

mattjj committed 29 days ago

acad3005

Merge pull request #35301 from mattjj:better-hijax-dce-interface

Google-ML-Automation committed 29 days ago

7f80f430

[hijax] improved dce interface

mattjj committed 29 days ago

b1981cbb

Canonicalize indexer.slice_sharding. Fixes https://github.com/jax-ml/jax/issues/35273

yashk2810 committed 29 days ago

7fbbe019

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

Google-ML-Automation committed 30 days ago

4f22f9d2

Merge pull request #35285 from mattjj:is-vjp-bit

Google-ML-Automation committed 30 days ago

4181fe49

[hijax] add is_vjp bit to LinearizeTrace

mattjj committed 30 days ago

79f3519e

Propagate input sharding to output in full_like if it's fully replicated. Then even if input/output shapes don't match, it should be fine. Fixes https://github.com/jax-ml/jax/issues/35273

yashk2810 committed 30 days ago

3b6fd04e

Bump actions/checkout from 6.0.0 to 6.0.1

dependabot[bot] committed 30 days ago

Verified 523b667b

Older