ROCm/jax

Pull Requests Commits

[ROCm] Simplify pytest parallelism calculation for ROCm CI

magaonka-amd committed 55 days ago

b0fb3495

Update XLA dependency to use revision

Google-ML-Automation committed 55 days ago

0c62bd53

Do not print the content of an array from `repr`/`str` if it's bigger than 1MiB

junwhanahn committed 56 days ago

f73e314d

Update XLA dependency to use revision

Google-ML-Automation committed 56 days ago

cec06d11

Merge pull request #39835 from mattjj:push-sqltwqqupsms

Google-ML-Automation committed 56 days ago

e2593fd2

Merge pull request #39837 from mattjj:push-yvqskzomlqvl

Google-ML-Automation committed 56 days ago

787b3a3e

[cleanup] remove process_call, now dead code

mattjj committed 56 days ago

82eef9fb

Remove the ifrt_version guard that's unnecessary

yashk2810 committed 57 days ago

73e3a60c

[remat3][vjp3] propagate zeros through remat3 via jax.vjp/linearize

mattjj committed 57 days ago

fea31ea9

Don't depend on `all_default_mem_kind` in device_put lowering

yashk2810 committed 57 days ago

57df5685

Update XLA dependency to use revision

Google-ML-Automation committed 57 days ago

8b9c8aba

[XLA:GPU]: Rolling back: Reattempt wrap all collectives in async-start for GPUs

yashk2810 committed 57 days ago

effc4526

Update XLA dependency to use revision

nitins17 committed 57 days ago

6c36869a

Merge pull request #39832 from mattjj:push-tpqzpvkzwpst

Google-ML-Automation committed 57 days ago

3be6807e

[refs] add pin=True/False arg to new_ref, delete config

mattjj committed 57 days ago

40fe4481

Merge pull request #39824 from mattjj:push-ukoxswropyol

Google-ML-Automation committed 57 days ago

6057d12a

Fix split pull rule sometimes propagating np ints.

Davis Yoshida committed 57 days ago

6e7cc852

revive lower_jaxpr_reference

mattjj committed 57 days ago

e52bc4cf

Update `is_matmul_supported` for TPU_8I + rephrase TPU_7 and TPU_7X.

Google-ML-Automation committed 57 days ago

7a92b495

Merge pull request #39788 from jax-ml:dependabot/github_actions/actions/cache/restore-6.1.0

Google-ML-Automation committed 57 days ago

7c1e21bb

Merge pull request #39819 from mattjj:push-znqzounwkztu

Google-ML-Automation committed 57 days ago

28341962

Merge pull request #39639 from Intel-tensorflow:oneapi_plugin_solver1

Google-ML-Automation committed 57 days ago

c3b5f855

[IFRT] Migrate IFRT RTTI users from llvm/mlir to xla::ifrt

hyeontaek committed 57 days ago

a82897d7

make leak-checker more robust in lower_jaxpr

mattjj committed 57 days ago

b668de20

[Mosaic TPU] Support packed types lane dim dynamic gather.

yueshengys committed 57 days ago

80c235cf

Merge pull request #39789 from jakevdp:broadcast-arrays-tuple

Google-ML-Automation committed 57 days ago

0fbedf7b

Merge pull request #39814 from fbourgey:typos_scipy_special

Google-ML-Automation committed 57 days ago

f07be2cf

Fix race condition in the python_callback_test.

nvgrw committed 57 days ago

591950d1

[triton] Added support for device-side TMA descriptors

superbobry committed 57 days ago

7df5ccc7

Bump actions/cache/restore from 4.3.0 to 6.1.0

dependabot[bot] committed 57 days ago

Verified cf8f53fe

Older