ROCm/jax

Pull Requests Commits

Release Prep JAX 0.9.1 (#735)

gulsumgudukbay committed 179 days ago

Verified 15f7ae2a

Prepare for JAX release 0.9.1

danielsuo committed 199 days ago

58cb6e55

Skip sparse random and dot general sampled ad tests.

danielsuo committed 199 days ago

bb3362eb

Skip layout_test.test_layout_donation_mismatching_in_and_out_fails for TPU v7 at libtpu<=0.3.5.

danielsuo committed 199 days ago

7aa10a02

Fix release workflow by using the folder with pre-downloaded jax wheels as artifact registry for libtpu wheels installation.

ybaturina committed 200 days ago

162b2ed0

Rename `pypi_latest` libtpu-version-type to `pinned` and use pinned libTPU version in corresponding workflows.

ybaturina committed 200 days ago

70a8149d

Skip tests that are failing at libtpu=0.0.35, but passing at libtpu-0.0.37.dev20260224.

danielsuo committed 200 days ago

0eb23146

[linux-arm64] Reduce parallelism to 32 for nogil pytest cpu to avoid OOMs.

danielsuo committed 200 days ago

969da1b7

Skip pallas lowering determinism testOrderAgnostic test if not on TPU.

danielsuo committed 200 days ago

a573e865

[jaxlib] Ported normalize_stubs.sh to Python

danielsuo committed 201 days ago

6428a91d

Add jaxlib version guard for updated make_c_api_client.

danielsuo committed 201 days ago

17cfbb81

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

danielsuo committed 201 days ago

c3b53a26

Roll back skipped tests that were failing on LLVM bit/byte-packing issue.

danielsuo committed 201 days ago

873bebe1

Suppress still-broken test.

brianwa84 committed 201 days ago

8f96afaa

Add a pre-commit hook that checks for Apache licence headers.

hawkinsp committed 201 days ago

66b61418

[JAX] Bound tracebacks up to the enclosing `jit`.

hawkinsp committed 201 days ago

6f2f0f36

Fix hijax `VmapOf.batch_dim_rule` with `None` map dims.

chr1sj0nes committed 201 days ago

6df7333b

Raise a better error if jit(shard_map) is bound with an AbstractMesh but called with an input with no mesh on it leading to num_devices mismatching.

yashk2810 committed 201 days ago

7af1783e

[Mosaic GPU] Relax restrictions on N for sparse tcgen05 MMA

apaszke committed 201 days ago

3f667fd4

[mosaic] Another go at removing tpu.wait_dma

superbobry committed 201 days ago

0b7d6672

Add int1/uint1 dtypes.

WindQAQ committed 201 days ago

c141c0fc

[Pallas:TPU] Add a test checking that dynamic shape exports require the pallas_export_experimental flag

apaszke committed 201 days ago

f2017f78

[Pallas:TPU] Add tests for DMAs with regular semaphores

apaszke committed 201 days ago

9e11f1c0

[Mosaic GPU] Add support for warp-level f8 MMAs

apaszke committed 201 days ago

7192cf7d

Merge pull request #35412 from levskaya:pyrefly

Google-ML-Automation committed 201 days ago

4826da47

Upgrade Abseil to LTS 20260107.1

akuegel committed 201 days ago

ee4cea11

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

Google-ML-Automation committed 201 days ago

d0d6f426

[Mosaic:GPU] Simplify collective metadata initialization.

PatriosTheGreat committed 201 days ago

6f06ab80

Reverts 6869e28a4bde0f6d3a06a7c39858fe40cc44ff60

Google-ML-Automation committed 202 days ago

efa6bd9f

Merge pull request #35414 from levskaya:pyrefly3

Google-ML-Automation committed 202 days ago

dc3f50ca

Older