Skip to content

Support dynamic expression clustering and broadcast singletons - #28

Open
stevenvar wants to merge 3 commits into
for-serving-2.20from
steven.tile-targeted-compile-time-guard
Open

Support dynamic expression clustering and broadcast singletons#28
stevenvar wants to merge 3 commits into
for-serving-2.20from
steven.tile-targeted-compile-time-guard

Conversation

@stevenvar

@stevenvar stevenvar commented Jul 22, 2026

Copy link
Copy Markdown
Owner

Summary

This PR improves how dynamic dimension expressions are handled when TensorFlow builds and runs XLA clusters.

It allows compatible expressions to share one cluster even when they have different constant offsets, while still rejecting expressions that depend on different variables. It also handles dimensions that may be runtime broadcast singletons without treating their observed value as strong evidence when solving the dynamic batch expression.

Changes

  • find the smallest sub-expression containing all variables and use it as the shared dynamic core
  • normalize that shared core to one fresh variable when compiling an XLA cluster
  • check cluster inputs and outputs for compatible variable sets and expression structure
  • preserve broadcast-singleton information from TensorFlow shape inference to XLA arguments
  • ignore conflicting singleton-derived solve candidates when consistent non-singleton candidates are available
  • add an optional, disabled-by-default majority-vote fallback for conflicting solve candidates
  • add tf_xla_disable_dynamic_size_padding to run dynamic sizes without padding them to a compile batch
  • return clear errors when dynamic expressions cannot be solved or substituted safely

Tests

  • compatible expressions such as (A + B) + 2 and (A + B) + 3 share the same clusterable core
  • expressions with different variable sets are rejected
  • single-variable expressions with different surrounding arithmetic remain compatible
  • the smallest sub-expression covering all variables is selected correctly
  • every matching occurrence of a shared sub-expression is replaced

This branch is based on the cleaned MIAC implementation and rebased onto the current for-serving-2.20.

@stevenvar
stevenvar force-pushed the steven.tile-targeted-compile-time-guard branch 8 times, most recently from 46c8c0b to d3fac93 Compare July 28, 2026 13:23
@stevenvar
stevenvar force-pushed the steven.tile-targeted-compile-time-guard branch from d8b93c9 to df4822c Compare July 29, 2026 17:19
@stevenvar stevenvar changed the title Steven.tile targeted compile time guard Support dynamic expression clustering and broadcast singletons Jul 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant