Commits ROCm/jax - SemanticDiff

Build jaxlib with the upstream Linux config

psanal35 committed 13 hours ago

0e8eef37

Drop address-of in EncapsulateFfiHandler calls

psanal35 committed 14 hours ago

ef298c70

Use jaxlib buffer callback handlers when available

psanal35 committed 14 hours ago

64096702

Build jax and jaxlib in the continuous wheel test

psanal35 committed 22 hours ago

99d615e6

Pass buffer callback handler pointers to bit_cast.

psanal35 committed 23 hours ago

df9d2617

Move buffer callback handlers into jaxlib

psanal35 committed 23 hours ago

26165d1f

Baseline before moving the buffer callback into jaxlib

psanal35 committed 23 hours ago

8422f851

[pallas:triton] Removed all tests

superbobry committed 1 day ago

a67386b6

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

sohaibiftikhar committed 1 day ago

07ca2ecb

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

golechwierowicz committed 1 day ago

fca2e320

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

golechwierowicz committed 1 day ago

df862f42

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

golechwierowicz committed 1 day ago

0c5186c9

[pallas] Moved `gpu_info` out of Pallas Triton

superbobry committed 1 day ago

88f525fc

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

Google-ML-Automation committed 1 day ago

d06a9ea0

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

Google-ML-Automation committed 1 day 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

jnp.nonzero: respect fill_value for empty arrays

Jake VanderPlas committed 1 day ago

0e3b4af9

Allow Format in jax.numpy.from_dlpack type annotations and docstring.

Google-ML-Automation committed 1 day ago

3197db94

Check for in_avals/out_avals.is_high too in pallas_call_p and mpmd_map_p because we strip mat from avals when we trace the kernel so jaxpr.is_high does not encapsulate the correct information when inputs/outputs are unreduced.

yashk2810 committed 1 day ago

5e34bd2f

`pl.Buffered` and `pltpu.BufferedRef.create` accepted `use_lookahead=True` even when `in_buffer_count` was less than two. Input lookahead prefetches future blocks into spare VMEM slots while the kernel body consumes the active slot. A single input slot leaves no spare slot for prefetching, and pure output buffers have zero input slots.

SamGinzburg committed 1 day ago

e8937d40

[jax-buildcop] Skip lax_numpy_operators_test on GPU ASAN and scipy_fft_test on CPU ASAN

Google-ML-Automation committed 1 day ago

b8fb21cc

Add unreduced/reduced rules for top_k and approx_top_k.

Angelogeb committed 1 day ago

ec3faf39

Implement `jax.scipy.special.yn` for integer-order Bessel functions of the second kind.

Jake VanderPlas committed 1 day ago

5d180d80