ROCm/jax

Pull Requests Commits

Only consider Tracers with ShapedArray avals as instances of Array.

hawkinsp committed 27 days ago

1b6e0981

[export] Remove serialization of old 16-bit nr_devices

gnecula committed 27 days ago

9f327f69

Fix compute_on2 with out_memory_spaces == Host

yashk2810 committed 27 days ago

191504b6

Include source info annotations in state primitive custom pretty printer rules.

hawkinsp committed 27 days ago

6dc73528

Merge pull request #34644 from ROCm:rocm-skip-prim-tri-solve-test

Google-ML-Automation committed 27 days ago

3a38a4bf

[pmap] Inline _prepare_pmap and clean up unused structs.

danielsuo committed 27 days ago

692d820f

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

Google-ML-Automation committed 27 days ago

762f0661

[pmap] Remove check_arg on every pmap invocation.

danielsuo committed 27 days ago

08fd5ea3

Switch XLA archive download URL from GitHub API to direct archive.

nitins17 committed 27 days ago

b839e6c3

Make `halt-for-connection` passable from workflow call.

ybaturina committed 27 days ago

44cc4f85

[ROCm] Added skip for "test_prim_tridiagonal_solve" tests on ROCm.

tsrw2048 committed 27 days ago

5eb5a215

[pmap] Add fast-paths for inferring and mapping in_axes.

danielsuo committed 27 days ago

d1fff7f5

[Mosaic GPU][NFC] Pass `CompiledKernel` by pointer to `CachedInit`.

allanrenucci committed 27 days ago

745cbe2f

[Pallas/Mosaic GPU] Use `self.capture_stdout` instead of `jtu.capture_stdout` in `mosaic_gpu_test`.

bchetioui committed 27 days ago

7241ee09

Fix XLA integrity in MODULE.bazel

meteorcloudy committed 27 days ago

9663c503

[Pallas] Disable a GPU interpreter test that deadlocks sometimes

apaszke committed 27 days ago

5cab99f9

[Mosaic GPU] Slightly refactor the argument idx search routine

apaszke committed 27 days ago

f0693d90

[Mosaic GPU] Generalize the fast i8 -> bf16 conversion to allow more vector lengths

apaszke committed 27 days ago

3b5d7209

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

Google-ML-Automation committed 28 days ago

de7e49e1

Merge pull request #34610 from magaonka-amd:skip-sparse-hipsparse-tests-upstream

Google-ML-Automation committed 28 days ago

cf73ceae

Correctly check dtype equality in `is_matmul_supported`.

Google-ML-Automation committed 28 days ago

979b537a

Support `device_put` of an uncommitted array to a single global device in McJAX.

emilyfertig committed 28 days ago

40b9554d

Reverts d3259a54160ee3a3724a17dce143d3212e5454b9

nitins17 committed 28 days ago

e0b62e82

Merge pull request #34632 from vfdev-5:fixd-typos-in-training-cookbook-

Google-ML-Automation committed 28 days ago

9f243da9

Merge pull request #34643 from mattjj:nodceeffect-hasheq

Google-ML-Automation committed 28 days ago

0d66c531

add __hash__ / __eq__ to NoDCEEffect for serialization.py

mattjj committed 28 days ago

778ec599

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

Google-ML-Automation committed 28 days ago

d3259a54

_pathways.pyi: Allow creating an XLA CPU client.

Google-ML-Automation committed 28 days ago

18a983dc

Skip sparse tests on ROCm due to hipSPARSE issue

magaonka-amd committed 28 days ago

3acc4108

[Mosaic GPU] Move kernel compilation out of the first execution.

allanrenucci committed 28 days ago

66e20704

Newer Older