ROCm/jax

Pull Requests Commits

[ROCm] Fix ROCm version in plugin wheel metadata

gulsumgudukbay committed 7 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 8 days ago

a7e4508c

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

rdyro committed 8 days ago

7c5346f2

Disable failing TPU v8i Pallas SC tests.

emilyfertig committed 8 days ago

d10c570e

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

rishipal committed 8 days ago

bb5d9b14

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

bchetioui committed 8 days ago

69aee06c

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

hawkinsp committed 8 days ago

41811a0c

Lower JAX_TEST_NUM_THREADS to 32 from 64 for TPU tests.

hawkinsp committed 8 days ago

6ef54da7

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

hawkinsp committed 8 days ago

823abec3

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

belitskiy committed 8 days ago

03a0cb24

Bump test_multiprocess time limits.

hawkinsp committed 8 days ago

2b07cfcc

[Mosaic GPU] Export GPU interpret mode from plgpu

Paul Biberstein committed 8 days ago

d217140a

Update XLA dependency to use revision

Google-ML-Automation committed 8 days ago

1507826a

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

Jake VanderPlas committed 8 days ago

2f06dd5b

[typing] Modernize types in jax.numpy

Jake VanderPlas committed 8 days ago

bed09490

Merge pull request #39981 from hawkinsp:libtpu2

Google-ML-Automation committed 8 days ago

5d4a37c8

Added cp.async support to the WG lowering

superbobry committed 8 days ago

18a18fb2

[pallas] Fixed the branching in :ops_test

superbobry committed 8 days ago

f96e3438

Regenerate requirements lock files.

hawkinsp committed 8 days ago

5fcd86f2

Fix Mac OS build failures after switching to bzlmod builds.

hawkinsp committed 8 days ago

889a9517

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

superbobry committed 8 days ago

a857d04f

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

superbobry committed 8 days ago

bf7de1fc

Merge pull request #39963 from jakevdp:test-types

Google-ML-Automation committed 8 days ago

f553ce05

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

Google-ML-Automation committed 8 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 8 days ago

dfb2a19c

Add a LIT (Layout In Types) API!

yashk2810 committed 9 days ago

cd036262

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

Google-ML-Automation committed 9 days ago

93374e8d

Merge pull request #39940 from mattjj:push-mupuxlypstot

Google-ML-Automation committed 9 days ago

eb19c69d

Bump actions/checkout from 6.0.3 to 7.0.1

dependabot[bot] committed 9 days ago

Verified 2e2cb348

[custom_vjp3] support closing over tracers via closure conversion

mattjj committed 9 days ago

91b72509

Older