ROCm/jax

Pull Requests Commits

Merge pull request #33701 from jakevdp:dtype-dep

Google-ML-Automation committed 36 days ago

a528b40d

improve error message and hasattr performance for Tracer.sharding

mattjj committed 36 days ago

c328f062

Add Pallas core_map guide

IvyZX committed 36 days ago

cdb2b692

Fix deprecated use of arrays in place of dtypes

jakevdp committed 36 days ago

c37f3904

Update hermetic CUDA UMD version to match CUDA driver version installed on RBE machines.

ybaturina committed 36 days ago

bcfea982

[Pallas] Fix cost estimate einsum bug.

justinjfu committed 36 days ago

eaf6efc1

PR #33677: [no-thunks] Reduce stack frames in `cond` and friends

dougalm committed 36 days ago

23cd412c

Merge pull request #33419 from Prakharprasun:doc-lax-conv-transpose

Google-ML-Automation committed 36 days ago

4d638b7e

Merge pull request #33688 from jakevdp:fix-dep-dtype

Google-ML-Automation committed 36 days ago

3b5c07c5

Merge pull request #33504 from mingxu1067:mingh/allow_256H_SDPA

Google-ML-Automation committed 36 days ago

4fad55b8

Build JAX artifacts only in CPU pools for RBE configurations.

ybaturina committed 36 days ago

8f48ac88

Fix deprecated use of arrays in place of dtypes

jakevdp committed 36 days ago

dd01b37c

[Mosaic] Support packed BroadcastInSublanesOp.

WindQAQ committed 36 days ago

307bca0d

Merge pull request #33685 from gnecula:export_fix_fwd_compat

Google-ML-Automation committed 36 days ago

6fa5ef52

[Pallas:MGPU] Fix Pallas `swap` lowering rule for scalar constants.

allanrenucci committed 36 days ago

b8dbb940

[export] Fix forward compatibility for exporting nr_devices

gnecula committed 36 days ago

746e13ef

Merge pull request #33632 from justinjfu:pallas_docstrings

Google-ML-Automation committed 36 days ago

e43b6cfe

[Mosaic GPU] Add support for reductions when writing to GMEM via TMA in WG semantics.

dimitar-asenov committed 36 days ago

403522de

Fix the jax line search, the zoom may only store two points in lo, hi and rec.

Google-ML-Automation committed 36 days ago

88d98bdb

Allow colocated python objects to use reference-counting to understand when the object can be destructed at the backend.

Google-ML-Automation committed 36 days ago

d23905a7

Export `reshard` as `jax.reshard`

yashk2810 committed 36 days ago

7b342dc9

Fix pjit_test on tpu v5_x4

yashk2810 committed 37 days ago

cafc0325

Update `rules_ml_toolchain` to disable `--@cuda_driver//:include_cuda_umd_libs` by default.

ybaturina committed 37 days ago

de608ee4

Fix footnote typo, move `jax.sharding.Mesh` to `jax.make_mesh`

ZacCranko committed 37 days ago

4592cfb3

Disable some JAX tests on TPU in open source.

jakeharmon8 committed 37 days ago

4d9ff5bd

Fix device_put when the destination mesh differs from context mesh under explicit mode

yashk2810 committed 37 days ago

dba8a62f

Merge pull request #33560 from jax-ml:ufunc-sigs

Google-ML-Automation committed 37 days ago

4ac5e1ec

Merge pull request #33067 from ROCm:fix/rocm7_plugin_dependency

Google-ML-Automation committed 37 days ago

f634fb27

Merge pull request #33674 from jakevdp:dep-concrete

Google-ML-Automation committed 37 days ago

913bc017

Merge pull request #33601 from jakevdp:dtype-protocol

Google-ML-Automation committed 37 days ago

b2ffd47b

Newer Older