ROCm/jax

Pull Requests Commits

[ROCm] Skipped a set of tests on ROCm that are flaky and sometimes cause CI failures.

tsrw2048 committed 108 days ago

c5d4aac6

Merge pull request #36225 from jakevdp:finalize-args

Google-ML-Automation committed 108 days ago

12786a78

Reverts d40a8e9d2dc4616806311ce8e91aafae73fdc46e

apaszke committed 108 days ago

6cdb59e0

[Pallas] Allow `pl.ds(size)` where size is a tracer. This is valid for dynamically sized DMAs.

brianwa84 committed 108 days ago

8558ea6c

Merge pull request #36238 from superbobry:pyrefly

Google-ML-Automation committed 108 days ago

25b709b7

Temporarily pinned jaxlib to the stable version in the Pyrefly pre-commit hook

superbobry committed 108 days ago

20299b72

Bump the minimum supported SciPy version for JAX to 1.14 and clean up obsolete version checks in tests.

hawkinsp committed 108 days ago

51d09f5e

[Mosaic GPU] Run more tests under WG semantics in distributed pallas test.

golechwierowicz committed 108 days ago

01860a05

[Mosaic GPU] Add WG lowering in Pallas for `try_cluster_cancel` and `query_cluster_cancel`

dimitar-asenov committed 108 days ago

d9c3946f

Another batch of Pyrefly-related fixes

superbobry committed 108 days ago

169ec187

[Mosaic GPU] Add layout inference and dialect lowering for `try_cluster_cancel` and `query_cluster_cancel`

dimitar-asenov committed 108 days ago

7784a1a6

Restore recently deleted skips (but disable them if libtpu is new)

apaszke committed 108 days ago

2b960779

[Mosaic] Enforce tpu.all_reduce of i1 has i32 result type

tlongeri committed 108 days ago

15e0086a

[mosaic] Force the minimum viable tiling for 1D buffers in infer-memref-layout

superbobry committed 108 days ago

2f8ff76d

[Mosaic GPU] Add support for block-scaled 2CTA MMAs with M=128 (64 per block)

apaszke committed 108 days ago

1e7b972e

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

Google-ML-Automation committed 108 days ago

1f0c269b

Finalize some deprecations for JAX v0.10.0:

Jake VanderPlas committed 108 days ago

2876211b

Keep the axis_env `axes` in mt.varying when remove_one_sized_mesh_axes is True

yashk2810 committed 109 days ago

311fc965

[Pallas TPU] Refactor Pallas Device ID Logic to use JAX lowering

sharadmv committed 109 days ago

88cd785d

Reverts b46adb327646222617e69c769436ec541ba91488

levskaya committed 109 days ago

6e5a5ca7

Merge pull request #36227 from mattjj:device-put-unreduced

Google-ML-Automation committed 109 days ago

dcad0e57

make device_put(x, P(unreduced={'x'})) work

mattjj committed 109 days ago

2f3557a3

propagate error to all sparse cores.

Google-ML-Automation committed 109 days ago

01032b33

lax.dot: remove custom TypeError for unexpected positional args

jakevdp committed 109 days ago

3cd96d00

Merge pull request #36186 from gnecula:consts3

Google-ML-Automation committed 109 days ago

4a010e0b

Add a manual pipeline test to SC

IvyZX committed 109 days ago

f9f84170

Update code that fails with NamedSharding(Mesh(device, ()), P()).

pschuh committed 109 days ago

c6ea0422

Merge pull request #36217 from jakevdp:default-dtype-bits

Google-ML-Automation committed 109 days ago

edf896a0

Add ctrl-c support for compilation now that there is a future pool:

pschuh committed 109 days ago

e7eec79c

Merge pull request #36216 from jakevdp:deprecation-removals

Google-ML-Automation committed 109 days ago

4ecf1632

Older