Skip to content

Preserve dynamic shape expressions through constant folding - #22

Merged
stevenvar merged 11 commits into
for-serving-2.20from
steven.partial_symbolic_const_folding
Jul 29, 2026
Merged

Preserve dynamic shape expressions through constant folding#22
stevenvar merged 11 commits into
for-serving-2.20from
steven.partial_symbolic_const_folding

Conversation

@stevenvar

@stevenvar stevenvar commented Apr 23, 2026

Copy link
Copy Markdown
Owner

Summary

This PR preserves dynamic shape expressions when TensorFlow partially evaluates shape-derived values and replaces them with constants. Without this metadata, XLA sees only the concrete value used during compilation and can lose the relationship with the dynamic input size.

The change:

  • tracks symbolic contents through supported shape-value operations, including Shape, Slice, StridedSlice, Gather, Reshape, Pack, ConcatV2, and Prod;
  • stores the symbolic contents on folded Const nodes and restores them when building XLA constants and compile-time arguments;
  • skips constant folding when the symbolic result cannot be reconstructed safely; and
  • avoids partial replacement of multi-output operations on non-CPU devices.

Tests

  • FoldShapeFromDynamicArgPreservesContents: keeps the expressions produced by Shape.
  • FoldSliceOfDynamicShapePreservesContents: keeps the selected dynamic dimension through Slice.
  • FoldGatherOfDynamicShapePreservesContents: keeps a gathered dynamic dimension in a scalar result.
  • FoldReshapeOfDynamicShapePreservesContents: keeps symbolic contents when a one-element vector is reshaped to a scalar.
  • FoldPackOfDynamicShapePreservesContents: keeps dynamic and static entries when they are packed into one shape tensor.
  • DoNotFoldUnsupportedDynamicContentsTransform: leaves unsupported transformations unfurled instead of dropping their expressions.
  • NoReplacePartialOutput: checks that non-CPU multi-output operations are not partially replaced.

@stevenvar
stevenvar force-pushed the steven.partial_symbolic_const_folding branch 5 times, most recently from 74663cc to 3cda438 Compare April 27, 2026 10:49
stevenvar pushed a commit that referenced this pull request Jul 7, 2026
resolve the issue for the xla batch matcher flag is set but invalid
@stevenvar
stevenvar force-pushed the steven.partial_symbolic_const_folding branch from b108667 to 0bd94c3 Compare July 29, 2026 13:47
@stevenvar stevenvar changed the title Steven.partial symbolic const folding Preserve dynamic shape expressions through constant folding 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