ROCm/jax

Pull Requests Commits

Update ad.backward_pass to support non-linear functions of constants.

dfm committed 351 days ago

b7ecfdfd

[Mosaic GPU] Add support for tiling stores to refs using small tiling

apaszke committed 351 days ago

85c6b6a1

Merge pull request #26950 from lockwo:Owen/add-pmap-typehint

Google-ML-Automation committed 351 days ago

de78d2cc

Dedup non-ref constants closed in cond branch functions.

danielsuo committed 351 days ago

e6db7a9d

Update XLA dependency to use revision

Google-ML-Automation committed 351 days ago

bf95bf49

wenscarl committed 351 days ago

ccbe9f7c

roofline: Support computing flops for unary ops.

zacmustin committed 351 days ago

8095d842

[Mosaic TPU] Support fp8 upcast to f32

bythew3i committed 351 days ago

ff4310f6

Auto complete specs in a sharding if aval.ndim > len(sharding.spec) with `None`. So that for a 2D input, P('data') continues to work.

yashk2810 committed 351 days ago

e9486920

Merge pull request #26927 from skye:merge_release

Google-ML-Automation committed 351 days ago

4cab1183

Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays.

Google-ML-Automation committed 351 days ago

cd7f03f2

Prevent tracer leaks in scipy.special.expn

jakevdp committed 351 days ago

b441b2b7

Open source TPU-friendly ragged paged attention kernel.

bythew3i committed 351 days ago

4b49c035

[Mosaic GPU] Fix `scf.ForOp` lowering to put lowered ops at the right place.

dimitar-asenov committed 352 days ago

5d64b3d2

Minor bug fixes in error checking

ayaka14732 committed 352 days ago

8c89da7c

Build JAX wheels instead of installing it from the source repository

nitins17 committed 352 days ago

623865fe

[pallas:mosaic_gpu] Added WG lowering rule for `lax.bitcast_convert_type_p`

superbobry committed 352 days ago

2a340193

[pallas] Add support for `pl.dot` with `int8` inputs.

chr1sj0nes committed 352 days ago

d6b97c20

Update XLA dependency to use revision

Google-ML-Automation committed 352 days ago

16bb9190

[Pallas/Mosaic GPU] Enable `ops_test` for Mosaic GPU.

bchetioui committed 352 days ago

fe577b5d

Just use `jit` as the string in error messages instead of `jit` and `pjit` based on resource_env. This is to start deprecating the need for `with mesh` and replace it with `use_mesh(mesh)`.

yashk2810 committed 352 days ago

a67ab9fa

Add a note about uneven sharding and with_sharding_constraint. Fixes https://github.com/jax-ml/jax/issues/26946

yashk2810 committed 352 days ago

ba5349f8

Set `USERPROFILE` for Windows builds to fix CI issue.

Google-ML-Automation committed 352 days ago

c16f37d8

[Pallas] Adds a simple dynamic race detector for TPU interpret mode.

jburnim committed 352 days ago

016b351f

Merge pull request #26952 from garymm:vmap-arg

Google-ML-Automation committed 352 days ago

8571ad9f

Fix build rule for free-threaded python builds.

Google-ML-Automation committed 352 days ago

0913cd75

vmap mismatch size error message: handle *args

garymm committed 352 days ago

69d66f66

Fix ambiguous cpu definition for JAX wheels.

Google-ML-Automation committed 352 days ago

3edc068f

add pmap axes hints

lockwo committed 352 days ago

3e4dc0d4

[Mosaic GPU] Remove support for large tiles on Blackwell

apaszke committed 352 days ago

8df00e26

Newer Older