ROCm/jax

Pull Requests Commits

Test ROCm libc++ with libunwind EH registration on

psanal35 committed 1 day ago

e8054422

Test ROCm libc++ with libunwind EH registration off

psanal35 committed 1 day ago

818910d5

Build the ROCm plugin with static libc++

psanal35 committed 1 day ago

5ab0dbdf

Integrate LLVM at llvm/llvm-project@efda9c9e4fd4

Google-ML-Automation committed 1 day ago

34beac2b

[jax-buildcop] Increase GPU/TPU shard_count and disable sanitizers on pjit_test and partitioning_test.

Google-ML-Automation committed 1 day ago

f6feb355

[Mosaic GPU] Tie the async_copy GMEM slice bounds check to oob_mode.

golechwierowicz committed 1 day ago

c467e1a5

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

Google-ML-Automation committed 1 day ago

222a71ad

Prefix private imports with underscores in jax.numpy type stubs

Google-ML-Automation committed 1 day ago

289b6505

[jax-buildcop] Skip cross-shard reduce_xor execution on GPU in test_reduce_prod_and_bitwise

Google-ML-Automation committed 1 day ago

74b807f7

Prefix private imports with underscores in jax.numpy type stubs

Jake VanderPlas committed 1 day ago

e0859795

[partitioning] Support sharded collapsed gather and take_along_axis out_sharding.

rdyro committed 1 day ago

73169151

[partitioning] Support sharded approx_top_k.

rdyro committed 1 day ago

57689943

[partitioning] Support sharded top_k.

rdyro committed 1 day ago

8f39a188

Run the TensorFlow side of jax2tf_multiprocess_test on CPU when JAX uses a separately loaded libtpu.so plugin.

Srishti Srivastava committed 1 day ago

17cf4934

[partitioning] Support sharded reduce_xor via pxor all-reduce.

rdyro committed 1 day ago

840f0194

[XLA:CPU] Lower log and log1p to YNNPACK SIMD bitcode kernels.

seantalts committed 1 day ago

d5862451

[jax-buildcop] Shard ragged_collective_test across 4 shards.

Google-ML-Automation committed 1 day ago

f9dcc9c8

Thread `out_sharding` through composite `jax.random` samplers

yashk2810 committed 1 day ago

c0f23197

Thread `out_sharding` through `random_gamma_p` and gamma-derived samplers

yashk2810 committed 1 day ago

15d8d3f0

Thread `out_sharding` natively through elementary `jax.random` samplers

levskaya committed 1 day ago

eeb73fe6

Support `out_sharding` in `random_bits_p`, `random_split_p`, and `iota_2x32_shape_p`

levskaya committed 1 day ago

9d34e951

Add `out_sharding` to `jnp.{quantile,nanquantile,nanpercentile}` and remove `auto_axes` from `jnp.percentile`.

levskaya committed 1 day ago

22e60356

Prevent Explicit mesh leaks in jit constant outputs and custom_vjp DCE

yashk2810 committed 1 day ago

f17de352

Plumb `reduction_order=` onto `lax.ragged_dot` and `lax.ragged_dot_general`

mattjj committed 1 day ago

57669800

[jax-buildcop] Speed up `test_planar_snake` in `mosaic_gpu_test` via `jit` + `lax.map`.

Google-ML-Automation committed 2 days ago

154baa05

jax: Fix TPU Pallas interpreter argument binding for closed-over scalar variables in mpmd_map

Google-ML-Automation committed 2 days ago

a02cc273

Add complex32 and bcomplex32 dtypes to JAX.

Jake VanderPlas committed 2 days ago

b48ff81c

Merge pull request #42430 from kratos0718:fix-choice-empty-shape

Google-ML-Automation committed 2 days ago

d1636437

[jax-buildcop] [CI] Enable `faulthandler_timeout=300` and `--durations=20` across non-TPU `pytest` CI scripts.

Google-ML-Automation committed 2 days ago

70a1e8bb

[Pallas] Push ref transforms into `pl.select_ref` branches.

rdyro committed 2 days ago

30096791

Older