ROCm/jax

Pull Requests Commits

Don't set PYTHONWARNINGS=error for tests that use TensorFlow.

hawkinsp committed 311 days ago

673a02d6

Update lax.linalg.qr primitives to use registration helper functions.

dfm committed 311 days ago

ed10003a

Disable tests/debug_info_test.py:test_vjp_of_jit test. Currently failing Python 3.13 GitHub CI tests.

danielsuo committed 311 days ago

87a7158f

Set the mesh of the sharding during broadcast in vmap so that we don't hit an error during canonicalization. This is similar to https://github.com/jax-ml/jax/commit/bcd4048dd5aa1ed8da39f6be88b4f1f3715f77ec

yashk2810 committed 311 days ago

66037d10

[CI] Enable workflow_dispatch for the continuous workflow

MichaelHudgins committed 311 days ago

c664a0cd

Merge pull request #26657 from gnecula:debug_info_test

Google-ML-Automation committed 311 days ago

cb8ff825

Merge pull request #26649 from jakevdp:wrapped-fun

Google-ML-Automation committed 311 days ago

237b7941

Update internal unop primitive helper to pass kwargs to dtype rule.

dfm committed 311 days ago

09325d92

Update lax.linalg.lu primitive to use registration helper functions.

dfm committed 311 days ago

126909b6

Start adding primitive registration helper functions to lax.linalg.

dfm committed 311 days ago

a981e1c4

Reverts 9ad9c3d3c2b6b5c3ea736af9b4d8c595537de93c

Google-ML-Automation committed 311 days ago

65526760

Merge pull request #26607 from dfm:ffi-example-dtype-dispatch

Google-ML-Automation committed 311 days ago

0e0a04fa

[better_errors] Fix a debug_info test, and expand the docstring for the helper function

gnecula committed 311 days ago

edf401d7

Add dtype dispatching to FFI example.

dfm committed 311 days ago

ca23b749

Update XLA dependency to use revision

Google-ML-Automation committed 311 days ago

0ef5eb06

Internal compatibility change

Google-ML-Automation committed 311 days ago

b510127a

Set the mesh of `tangent.aval` when we are creating `zeros_like_aval` because when you close over an array which is unused, we error out during canonicalization

yashk2810 committed 312 days ago

bcd4048d

Use the mesh of `out_aval` when converting GSPMDSharding to NamedSharding. This makes sure that the axis types of the corresponding output is correct.

yashk2810 committed 312 days ago

250e2ee7

Merge pull request #26646 from jakevdp:fix-rtd

Google-ML-Automation committed 312 days ago

bbc4fa71

Internal: improved type annotations for lu.WrappedFun

jakevdp committed 312 days ago

fe00aa0f

[pallas:mosaic_gpu] Added support for binary/comparison ops with WG semantics

superbobry committed 312 days ago

7438976e

Fix head comment: was referring to nonexistent parameters.

lrdxgm committed 312 days ago

08de0128

CI: skip execution of convolutions.ipynb

jakevdp committed 312 days ago

380fb0a5

Fix wheel downloads for nightly artifact testing

nitins17 committed 312 days ago

6b39eb6b

[JAX] Generate more readable error for failed device deserialization in colocated Python

hyeontaek committed 312 days ago

71f9764e

[Pallas][Mosaic] Support float8_e4m3b11fnuz

Google-ML-Automation committed 312 days ago

b7968474

Added `jax.experimental.multihost_utils.live_devices` API.

mwhittaker committed 312 days ago

ddcb7dee

Bump up timeout for building wheel to 60 minutes

nitins17 committed 312 days ago

040b718e

Update XLA dependency to use revision

Google-ML-Automation committed 312 days ago

2c8a8ed7

[Mosaic] Consider divisibility when doing large tiling

Google-ML-Automation committed 313 days ago

37af0135

Newer Older