ROCm/jax

Pull Requests Commits

Build the ROCm plugin with static libc++

psanal35 committed 6 hours ago

17b2b528

Merge pull request #42195 from Woojay44:fix-dpa-residual-docstring

Google-ML-Automation committed 8 hours ago

d613257d

[Mosaic] Keep the reduced dimension in tpu.reduce_index

naummo committed 8 hours ago

ce319ea3

[Pallas/Mosaic GPU] Allocate SMEM refs using logical shapes in run_scoped.

allanrenucci committed 8 hours ago

f03b0382

Support `jax.shard_map` inside `pallas.fuser.fuse`

IvyZX committed 9 hours ago

d9cef03f

Merge pull request #42063 from BlazeSR10hawk:fix/scipy-signal-mixed-size-42061

Google-ML-Automation committed 9 hours ago

f1d5dc7d

Fix mixed-size scipy.signal convolutions

BlazeSR10hawk committed 10 hours ago

85892298

Avoid converting real inputs to complex in +, -, *, and complex / real

hawkinsp committed 11 hours ago

7c579666

Use jax.lax.sinpi, cospi, and tanpi in internal callers.

hawkinsp committed 11 hours ago

520f5a23

[pallas] Migrated the remaining few tests off the deprecated `pl.core_map`

superbobry committed 11 hours ago

5015b408

[pallas:triton] Removed the Triton backend

superbobry committed 12 hours ago

2361d533

Avoid overflow and reduce op count in the JVP rule for lax.div.

hawkinsp committed 14 hours ago

ce60b242

[pallas:triton] Removed all mentions of the Triton backend from the docs

superbobry committed 14 hours ago

748762f0

[jax-buildcop] [jax2tf] Relax CPU FFT numerical tolerance in jax2tf_limitations

Google-ML-Automation committed 14 hours ago

d603f306

[pallas:triton] Removed all tests

superbobry committed 15 hours ago

a67386b6

[JAX] Copy kernel_c_api.h and kernel_api.h in build_wheel.py.

sohaibiftikhar committed 17 hours ago

07ca2ecb

[Mosaic:GPU][NFC] Clarify why we skip certain types during conversions tests.

golechwierowicz committed 18 hours ago

fca2e320

[Mosaic GPU] Allow F8E5M3FNU scales in dialect ops and layout inference.

golechwierowicz committed 18 hours ago

df862f42

[Mosaic GPU] Add support for Float8E5M3FNU conversions and abs in FragmentedArray.

golechwierowicz committed 19 hours ago

0c5186c9

[pallas] Moved `gpu_info` out of Pallas Triton

superbobry committed 19 hours ago

88f525fc

[jax-buildcop] Disable `tpu_pallas_interpret_test` under ASAN where it times out

Google-ML-Automation committed 20 hours ago

d06a9ea0

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

Google-ML-Automation committed 20 hours ago

cb21c380

[jax-buildcop] Fix PallasHiJaxTest.test_pallas_call_shard_map_unreduced sharding spec and reference matmul.

Google-ML-Automation committed 1 day ago

0f19df23

Add jax.lax.sinpi, cospi, and tanpi.

hawkinsp committed 1 day ago

a024fff0

[partitioning] Add `p{max,min}_scatter` and lower `reduce_{max,min}` scatters.

rdyro committed 1 day ago

1d8b58a5

[unreduced] Strip `mat` from Refs in mpmd map unconditionally.

rdyro committed 1 day ago

fbbd757c

Support 64-bit integer BLAS and LAPACK when JAX is linked directly against LAPACK.

hawkinsp committed 1 day ago

f206f40f

Remove the `order_wrt_mesh(_get_spec_set) pattern. This pattern has a potential to introduce bugs and should never be used. Also eemove _get_spec_set and introduce flatten_spec_no_none which returns a tuple. The callers can convert to a set if needed

yashk2810 committed 1 day ago

a8572377

[nn] Enable jax_softmax_custom_jvp=True by default

Jake VanderPlas CA committed 1 day ago

c324243a

[partitioning] Lower naryop elementwise ops to explicit shard_map rules.

rdyro committed 1 day ago

b695bbf9

Older