ROCm/jax

Pull Requests Commits

[ROCm] Fix ROCm version in plugin wheel metadata

gulsumgudukbay committed 52 days ago

6c26377d

Provide native convolution support (`pltpu.conv` and `lax.conv_general_dilated`) in JAX Pallas and the MLIR Mosaic TPU dialect for basic N-D spatial convolutions.

blakehechtman committed 52 days ago

a7e4508c

[fuser] Fix support for single-output stateful fusions.

rdyro committed 52 days ago

7c5346f2

Disable failing TPU v8i Pallas SC tests.

emilyfertig committed 52 days ago

d10c570e

Delete the floating point optimization level and memory fitting level settings in favor of the EffortLevel enum.

rishipal committed 52 days ago

bb5d9b14

[Pallas/Mosaic GPU] Skip tests using `cp.async` with warpgroup semantics when the op is not available.

bchetioui committed 52 days ago

69aee06c

Add Python 3.15 and 3.15t to the nightly/release test matrix.

hawkinsp committed 52 days ago

41811a0c

Lower JAX_TEST_NUM_THREADS to 32 from 64 for TPU tests.

hawkinsp committed 52 days ago

6ef54da7

Second attempt at fixing Mosaic GPU CI failures after bzlmod switch.

hawkinsp committed 52 days ago

823abec3

[CI] Patch rules_python to pass extended-length paths even on newer Windows.

belitskiy committed 52 days ago

03a0cb24

Bump test_multiprocess time limits.

hawkinsp committed 52 days ago

2b07cfcc

[Mosaic GPU] Export GPU interpret mode from plgpu

Paul Biberstein committed 52 days ago

d217140a

Update XLA dependency to use revision

Google-ML-Automation committed 52 days ago

1507826a

[typing] upgrade uses of typing.Union and typing.Optional

Jake VanderPlas committed 52 days ago

2f06dd5b

[typing] Modernize types in jax.numpy

Jake VanderPlas committed 52 days ago

bed09490

Merge pull request #39981 from hawkinsp:libtpu2

Google-ML-Automation committed 52 days ago

5d4a37c8

Added cp.async support to the WG lowering

superbobry committed 52 days ago

18a18fb2

[pallas] Fixed the branching in :ops_test

superbobry committed 52 days ago

f96e3438

Regenerate requirements lock files.

hawkinsp committed 52 days ago

5fcd86f2

Fix Mac OS build failures after switching to bzlmod builds.

hawkinsp committed 52 days ago

889a9517

[pallas] Do not use `pl.pallas_call` in :export_back_compat_pallas_test for Mosaic GPU

superbobry committed 52 days ago

a857d04f

[pallas:mgpu] Do not use `pallas_call` in tests

superbobry committed 52 days ago

bf7de1fc

Merge pull request #39963 from jakevdp:test-types

Google-ML-Automation committed 52 days ago

f553ce05

Merge pull request #39475 from edluyuan:fix/nanquantile-empty-reduction

Google-ML-Automation committed 53 days ago

664e8494

[Mosaic TPU] Support reshapes of form `[m, n] -> [m*n, 1]` and `[m*n] -> [m, n, 1]` where `n % 128 == 0`.

yueshengys committed 53 days ago

dfb2a19c

Add a LIT (Layout In Types) API!

yashk2810 committed 53 days ago

cd036262

Merge pull request #39964 from jax-ml:dependabot/github_actions/actions/checkout-7.0.1

Google-ML-Automation committed 53 days ago

93374e8d

Merge pull request #39940 from mattjj:push-mupuxlypstot

Google-ML-Automation committed 53 days ago

eb19c69d

Bump actions/checkout from 6.0.3 to 7.0.1

dependabot[bot] committed 53 days ago

Verified 2e2cb348

[custom_vjp3] support closing over tracers via closure conversion

mattjj committed 53 days ago

91b72509

Older