ROCm/jax

Pull Requests Commits

[mosaic] Propagate the layout from tpu.erase_memref_layout operand to tpu.log_buffer

superbobry committed 140 days ago

9c25bbc2

[Mosaic GPU] Lock down 8-bit pointwise ops (except bitwise and cmp) + add lshift/rshift

apaszke committed 140 days ago

0cd5ca77

[Mosaic GPU] Reenable the use of redux.sync, but only when all lanes participate

apaszke committed 140 days ago

51d35912

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

Google-ML-Automation committed 141 days ago

7a5e13d7

Add Fuser support for jnp.tile input fusions

Google-ML-Automation committed 141 days ago

156f7d2e

Bump XLA version.

mwhittaker committed 141 days ago

cdfe904a

Add `testonly=True` to py_import targets that depend on testonly `wheel_additives` target.

ybaturina committed 141 days ago

a84ca0c5

[Pallas] Create trace_value primitive for dynamic value logging

Google-ML-Automation committed 141 days ago

02939fef

Merge pull request #34418 from jakevdp:hijax-shard-map

Google-ML-Automation committed 141 days ago

6a9f3d98

Follow up pr after triton integration cl and tokamax pr.

loislo committed 141 days ago

91223b93

Migrate TSAN workflow to use RBE.

belitskiy committed 141 days ago

acdd012a

Merge pull request #34410 from jakevdp:indexing-strides

Google-ML-Automation committed 141 days ago

f0ff1b85

Merge pull request #34414 from jakevdp:padtype-to-pads

Google-ML-Automation committed 141 days ago

1143012a

[pmap] Suppress PmapSharding deprecation warning in (more) multihost_utils_test.

danielsuo committed 141 days ago

3d400948

[hijax] fix shard_map of hijax primitive

jakevdp committed 141 days ago

2ecafeee

[Pallas] Disable non-sharded multi-chip splash attention test config, since the additional chips are unused there.

rdyro committed 141 days ago

9e481477

Uninstall xprof on python3.13-nogil always.

mwhittaker committed 141 days ago

fe7676d4

Merge pull request #34394 from mattjj:custom-vjp3-youandme-4

Google-ML-Automation committed 141 days ago

429fcfc0

[typ] annotate lax.padtype_to_pads

jakevdp committed 141 days ago

e64dccf8

[indexing] use strides=None for all unit strides

jakevdp committed 141 days ago

90076305

[Pallas] Remove PRNG/Semaphore effects from basic JAX allowlist since

sharadmv committed 141 days ago

ee85b82f

[Pallas] Turn off tpu_7x configs for distributed test to avoid hangs

sharadmv committed 141 days ago

7f6fe293

[hijax] fix up custom_vjp3 error messages

mattjj committed 141 days ago

002c3dc3

Merge pull request #34375 from Cjkkkk:relax_fp8_sdpa_tolerance

Google-ML-Automation committed 141 days ago

c5831d16

[Pallas:MGPU] Fix scratch size for cross-warp reductions.

allanrenucci committed 141 days ago

695eac8e

Add the test rule `compare_srcs_and_test_deps_test` that compares the tests' dependencies and wheel sources to validate that the wheels have all required files.

ybaturina committed 141 days ago

8d97179e

Add a version guard around TPU lowering rule changes.

hawkinsp committed 141 days ago

e16ca653

Update `rules_ml_toolchain` version to incorporate new GPU folders structure and add hermetic ASAN support.

ybaturina committed 141 days ago

a8927894

[pmap] Deprecate setting `jax_pmap_shmap_merge` config state.

danielsuo committed 141 days ago

bb1017b8

[pallas:sc] The lowering can now use an empty grid directly

superbobry committed 141 days ago

0c804685

Newer Older