Commits ROCm/jax - SemanticDiff

Enable remote compilation for jax in rocm rbe

alekstheod committed 3 days ago

2cf93adb

[jax] Simplify test_closed_over_constant_sharded_and_cache_hit in pjit_test

gnecula committed 4 days ago

097222b8

Document possible migration failures when we turn on use_simplified_jaxpr_constants.

gnecula committed 4 days ago

81450934

[Pallas:TPU] Return final loop index carry in scan and while fori lowering

SamGinzburg committed 4 days ago

f7d9eefd

[Pallas:SC] Flip sign bit for signed int32 keys in SparseCore lax.sort lowering

SamGinzburg committed 4 days ago

4460acc2

raise NotImplementedError for primitives that don't support layout rules yet and not implement themselves via standard_primitive

yashk2810 committed 4 days ago

843b82a5

Added layout rules and lower_with_explicit_types lowerings for integer_pow_p, reduce_precision_p, split_p, clamp_p, concatenate_p, select_n_p, slice_p, dynamic_slice_p, and dynamic_update_slice_p

yashk2810 committed 4 days ago

e6534d02

Fall back to caller location for bare-NameLoc MLIR callsites in source_mapper.

levskaya committed 4 days ago

11def02c

Merge pull request #41768 from bhaskargurram-ai:perf/accumulate-associative-scan

Google-ML-Automation committed 4 days ago

5ded87b6

Merge pull request #41858 from ayaka14732:pallas-sc-fstring

Google-ML-Automation committed 4 days ago

13e77957

[Mosaic GPU] Guard `bf16` <-> `f8e4m3fn`/`f8e5m2` conversions on Blackwell+ (`sm_100`+).

Google-ML-Automation committed 4 days ago

62045164

[emit_pipeline] Support PyTree window views, allocations, scratches, and private window effects in emit_pipeline primitive

rdyro committed 4 days ago

d330fa6d

Merge pull request #41856 from stroxler:fix-noncontiguous-overload

Google-ML-Automation committed 4 days ago

a320c518

[Pallas SC] Add missing f-string prefixes in get/swap lowering error messages

ayaka14732 committed 4 days ago

03fec645

Lower lax.exp2_p and lax.log2_p directly to stablehlo.exp2 and stablehlo.log2.

hawkinsp committed 4 days ago

eaf0cbf3

[typing] Avoid non-contiguous overloads in splash_attention_kernel.py

stroxler committed 4 days ago

f6d00a4b

Add worst-case math library inputs to JAX numerics accuracy tests.

hawkinsp committed 4 days ago

b7fab121

[Mosaic GPU] Replace ScopedMockLog with a log sink in custom_call_test.

allanrenucci committed 4 days ago

7016a4de

[jax2tf] Revert Shardy input reshard count expectation in `test_jit_sharding_constraint`.

Google-ML-Automation committed 4 days ago

028d0db5

Merge pull request #41230 from ROCm:pemeliya/mori-deps-rpath-fixes

Google-ML-Automation committed 4 days ago

889ca6a2

[Mosaic GPU] Support `.rz` rounding for narrowing FP8 casts on `sm_107`.

bchetioui committed 4 days ago

e7105a80

Mark `lax_test` as `notsan` on CPU due to TSan timeouts.

Google-ML-Automation committed 4 days ago

0809c940

[mgpu] Simplify `utils.inline_ptx` signature and clarify `has_side_effects` default.

chr1sj0nes committed 4 days ago

581511e1

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

Google-ML-Automation committed 5 days ago

750ff493

[serialize] Add serialization versioning to serialize_executable

gnecula committed 5 days ago

4fc39d09

Increase barrier and key-value sync timeout in `WaitBarrierTest` from 1s to 3s.

Google-ML-Automation committed 5 days ago

72b2dd70

Remove unused `std::string` variable to fix `-Wunused-variable` warning.

nigham committed 5 days ago

d6b5301f

Mark `lax_control_flow_test`, `mutable_array_test`, and `lax_numpy_reducers_test` as `noasan` due to ASan timeouts.

Google-ML-Automation committed 5 days ago

bda56689

[Pallas:SparseCore]Fix SPMEM-to-SC SMEM return-before-copy

Jake VanderPlas committed 5 days ago

1937fa70

[mgpu] Lower two-branch conditionals to scf.if.

chr1sj0nes committed 5 days ago

b13f30b1