ROCm/jax

Pull Requests Commits

[ROCm] Skip `tridiagonal_solve_perturbed` code path in eigh unit tests on ROCm.

tsrw2048 committed 91 days ago

698c52c3

Merge pull request #36936 from jakevdp:key-bool

Google-ML-Automation committed 91 days ago

de15bbeb

[Pallas SC] Support indexers that have both indirect offsets and minor slices

apaszke committed 91 days ago

5112029a

[Pallas/Mosaic GPU] Fix lowering of `_device_id_to_logical`.

bchetioui committed 91 days ago

83c8ce4a

[Pallas SC] Disallow gathers with 2D+ indices

apaszke committed 91 days ago

489e6770

Update XLA dependency to use revision http://github.com/openxla/xla/commit/8b5fd56745d1deaacee5a142838b80398a04dbc9

Google-ML-Automation committed 92 days ago

fa4a5550

[Mosaic GPU] Clean up guards for min jaxlib version.

golechwierowicz committed 92 days ago

9b21959f

Update XLA dependency to use revision http://github.com/openxla/xla/commit/400cffe3ad2a5e8384ae2064bb454a02d1197ebc

Google-ML-Automation committed 93 days ago

a6e4a8b9

[Mosaic GPU] Don't reuse variable to encode layout equality in inference rules.

allanrenucci committed 93 days ago

fde69055

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

Google-ML-Automation committed 94 days ago

126813ee

[Pallas TPU] Trace with a DEFAULT memory space for pallas_call instead of None to disambiguate with ANY

sharadmv committed 94 days ago

833c5b78

Propagate memory space through physical aval

sharadmv committed 94 days ago

362232ff

Merge pull request #36887 from hawkinsp:reqs

Google-ML-Automation committed 94 days ago

bacb431d

[Reland] Add SparseCore availability checks during Pallas lowering.

Google-ML-Automation committed 94 days ago

5cfa48df

Add a faster path for core.is_literalable for scalar types.

hawkinsp committed 94 days ago

3c2d82a6

Merge pull request #36942 from jakevdp:type-promotion-error

Google-ML-Automation committed 94 days ago

ea458338

Serialize pin backend_config to JSON.

sparc1998 committed 94 days ago

edb184a0

Export jax.dtypes.TypePromotionError

jakevdp committed 94 days ago

a3fba0ec

[XLA SC] Automatically annotate all functions in Mosaic modules as transform indices.

pkanwar23 committed 94 days ago

e20fdb3e

[Colocated Python] Temporarily work around an incorrect CPU device setup in JAX tests

hyeontaek committed 94 days ago

08902d39

Replace deprecated `with mesh` with `with jax.set_mesh(mesh)` in tests

yashk2810 committed 94 days ago

c8dbbd11

[Pallas/Mosaic GPU] Add warp-level support for `semaphore_wait` with warpgroup semantics.

bchetioui committed 94 days ago

b3efe4e7

Explicitly define __bool__ for PRNGKeyArray

jakevdp committed 94 days ago

68706155

[Pallas/Mosaic GPU][NFC] Remove warpgroup-semantics transform specialization in `gpu_pallas_distributed_test.py`.

bchetioui committed 94 days ago

94074243

[Pallas/Mosaic GPU] Enable a few tests previously gated by user transform support.

bchetioui committed 94 days ago

5ddcfed4

[Pallas/Mosaic GPU] Carry `BlockSpec` transforms when lowering with warpgroup semantics.

bchetioui committed 94 days ago

f2d16874

Merge pull request #36919 from ezhulenev:patch-2

Google-ML-Automation committed 94 days ago

60e9ffac

[pallas:mgpu] Add `barrier_test` function.

chr1sj0nes committed 94 days ago

a7829283

[pallas:mgpu] Relax `dtype` type bounds.

chr1sj0nes committed 94 days ago

00f7e548

Implement AREA image resizing.

hawkinsp committed 94 days ago

dd42d820

Older