ROCm/jax

Pull Requests Commits

Follow up pr after triton integration cl and tokamax pr.

loislo committed 20 days ago

91223b93

Migrate TSAN workflow to use RBE.

belitskiy committed 20 days ago

acdd012a

Merge pull request #34410 from jakevdp:indexing-strides

Google-ML-Automation committed 20 days ago

f0ff1b85

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

Google-ML-Automation committed 20 days ago

1143012a

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

danielsuo committed 20 days ago

3d400948

[hijax] fix shard_map of hijax primitive

jakevdp committed 20 days ago

2ecafeee

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

rdyro committed 20 days ago

9e481477

Uninstall xprof on python3.13-nogil always.

mwhittaker committed 20 days ago

fe7676d4

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

Google-ML-Automation committed 20 days ago

429fcfc0

[typ] annotate lax.padtype_to_pads

jakevdp committed 20 days ago

e64dccf8

[indexing] use strides=None for all unit strides

jakevdp committed 20 days ago

90076305

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

sharadmv committed 20 days ago

ee85b82f

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

sharadmv committed 20 days ago

7f6fe293

[hijax] fix up custom_vjp3 error messages

mattjj committed 20 days ago

002c3dc3

Merge pull request #34375 from Cjkkkk:relax_fp8_sdpa_tolerance

Google-ML-Automation committed 20 days ago

c5831d16

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

allanrenucci committed 20 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 20 days ago

8d97179e

Add a version guard around TPU lowering rule changes.

hawkinsp committed 20 days ago

e16ca653

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

ybaturina committed 20 days ago

a8927894

[pmap] Deprecate setting `jax_pmap_shmap_merge` config state.

danielsuo committed 20 days ago

bb1017b8

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

superbobry committed 20 days ago

0c804685

Merge pull request #34400 from gnecula:shape_poly_pallas_cache_reuse

Google-ML-Automation committed 20 days ago

d66c2a69

Merge pull request #34403 from gnecula:export_fix1

Google-ML-Automation committed 20 days ago

31ef45ef

[export] Cleanup the handling of has_named_sharding

gnecula committed 20 days ago

2df10cb3

[pallas] Make the LoweringDynamicShapeEnv use local mappings

gnecula committed 20 days ago

602f212d

Reverts 7cd6b1a9c711068f4f05312d24ef598fe40d4e50

Google-ML-Automation committed 21 days ago

15a3e1ed

[XLA:MGPU] Port Partitioned.*Dims and VectorLength methods to C++.

golechwierowicz committed 21 days ago

92f08100

[Mosaic GPU] Disable redux ops until properly benchmarked and optimized

Google-ML-Automation committed 21 days ago

ef84d389

[XLA:MGPU] Port TiledLayout's construction logic to C++.

golechwierowicz committed 21 days ago

9ac60a34

Merge pull request #33853 from gnecula:export_v6_tests

Google-ML-Automation committed 21 days ago

bcc959b9

Newer Older