ROCm/jax

Pull Requests Commits

Merge pull request #34943 from jurasic-pf:jurasic/faster-trapezoid

Google-ML-Automation committed 58 days ago

d09caeae

Merge pull request #35138 from jakevdp:pyrefly-fixes

Google-ML-Automation committed 58 days ago

985a077c

[checkify] delete non-hermetic test

mattjj committed 58 days ago

3d197e77

Use h100x2_full for the NVSHMEM test

apaszke committed 58 days ago

4aca5953

[Pallas] Use NumPy CPU random number generators instead of generating on TPU.

WindQAQ committed 58 days ago

d962a2bb

[Mosaic GPU] Add support for st.async (through FragmentedArray.store_tiled_async)

apaszke committed 58 days ago

115ceae8

[pallas:sc] The only supported array shapes now match physical vector register shapes

superbobry committed 58 days ago

26c69f1c

[Mosaic:GPU] Fix the cross-gpu barrier signal buffer sizes

PatriosTheGreat committed 58 days ago

6e6f672b

Google-ML-Automation committed 58 days ago

4175a8e3

[Mosaic GPU] Add support for async copies of arbitrary data from SMEM to TMEM.

apaszke committed 59 days ago

284864b3

[pallas:mosaic] Removed an unnecessary assignment to `MemorySpace.ANY`

superbobry committed 59 days ago

7aa83317

[Mosaic GPU] Slightly simplify the extract_strided_slice pattern

apaszke committed 59 days ago

0526a35c

[pallas] Merge the shape polymorphism tests into pallas_shape_poly_test.py.

gnecula committed 59 days ago

e8cad830

[jaxlib] Build MLIR .pyis and include them into the wheel

superbobry committed 59 days ago

c88e3b00

[Mosaic GPU] Add support for reductions of arrays that have untiled dimensions

apaszke committed 59 days ago

3e8cb42e

[Mosaic GPU][NFC] Move cluster-related functions out of lowering context

apaszke committed 59 days ago

5c78b63f

[Mosaic GPU] Add an optimized implementation for 1/x for x of type ue8m0

apaszke committed 59 days ago

2b21754f

[pallas:sc] Use an empty grid in `ScalarSubcoreTest.test_pallas_call`

superbobry committed 59 days ago

62e36c1e

[mosaic] Add optional `core_type` parameter to `tpu.memory_space`

superbobry committed 59 days ago

48080782

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

Google-ML-Automation committed 59 days ago

958b7869

Merge pull request #35115 from ROCm:rocm-hip-threefry-ffi-stable

Google-ML-Automation committed 59 days ago

e50a76c4

Add `remove_size_one_mesh_axis_from_type` to JaxprEqnContext so that it's preserved through jaxpr roundtrip

yashk2810 committed 59 days ago

c51913c5

Update ravel_multi_index to accept optional dtype input, and raise error if (possibly inferred) dtype is too small given the dims.

Ryan McKenna committed 59 days ago

32081044

Remove StatusOr from hlo_modules return type for compatibility with nanobind.

emilyfertig committed 59 days ago

e01608d3

jnp.trapezoid fastpath when uniform dx is provided

jurasic-pf committed 59 days ago

74aab53f

Move `experimental_compile_mpmd` from pathways to jaxlib.

emilyfertig committed 59 days ago

27274db2

Non-splat is unchanged

ImanHosseini committed 59 days ago

42c458cc

[Pallas] Shrink einshape generation size to avoid ASAN test timeouts

sharadmv committed 59 days ago

796d3bf8

[pyrefly] more typing fixes

jakevdp committed 59 days ago

ade6b64e

Fix `_broadcast_in_dim_transpose_rule` to only consider unit_dims if the shape == 1 and spec is None. Fixes https://github.com/jax-ml/jax/issues/35181

yashk2810 committed 59 days ago

e69ef7b5

Newer Older