Allow partial symbolic constant folding by stevenvar · Pull Request #98 · joeyye-work/tensorflow

stevenvar · 2026-04-23T11:25:47Z

This PR makes common-runtime constant folding less conservative for symbolic shape-derived values.

Before this change, TensorFlow often had to skip folding small shape-tensor subgraphs because replacing them with a plain Const would lose the symbolic contents needed later by tf2xla/XLA. This PR allows those folds when the symbolic contents can be preserved exactly, and still skips them when they cannot.

What’s included:

preserve symbolic contents for folded shape-tensor plumbing ops such as Shape, Slice, Gather/GatherV2, Reshape, Pack, ConcatV2, and supported Prod cases
store that symbolic contents metadata losslessly on replacement Const nodes using a serialized attr
add focused tests for Shape(arg), Slice(Shape(arg)), GatherV2(Shape(arg), 0, axis=0), and Reshape(Slice(Shape(arg)), [])

Copilot

Pull request overview

This PR relaxes TensorFlow constant-folding restrictions for shape-derived/symbolic values by allowing folding when the folded result can preserve exact symbolic “contents”, and propagates that preserved metadata through to XLA.

Changes:

Add symbolic-contents preservation logic in constant folding (new user_inferred_value_contents attr, plus a shape-tensor op subset that can be traced exactly).
Gate replacement of folded constants on whether symbolic contents can be preserved, rather than blocking folding up-front.
Update XLA constant lowering paths to consume user_inferred_value_contents, and add a regression test for Gather-from-shape.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

File	Description
tensorflow/core/common_runtime/constant_folding.cc	Adds symbolic-contents reconstruction (`TryGetFoldedValueContents`) and attaches `user_inferred_value_contents`; moves the symbolic safety guard to replacement time.
tensorflow/core/common_runtime/constant_folding_test.cc	Adds a test ensuring a folded Gather-from-Shape preserves selected symbolic contents.
tensorflow/compiler/tf2xla/kernels/const_op.cc	Allows ConstOp to attach contents from `user_inferred_value_contents` (or fallback to `user_inferred_shape`) for scalar/vector cases.
tensorflow/compiler/jit/kernels/xla_ops.cc	Allows XlaCompileOp to recover constant-value expressions from `user_inferred_value_contents` in addition to existing dynamic-shape metadata.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-27T11:41:08Z

+  TensorShapeProto exact_contents;
+  const bool has_exact_contents = TryGetFoldedValueContents(n, 0, &exact_contents);
+  const bool has_dynamic =
+      GetShapeFromDirectDynamicSource(n, &exact_contents);
+  const bool is_shape_derived =
+      n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr;
+  if ((has_dynamic || is_shape_derived) && (!has_exact_contents || n->num_outputs() > 1)) {
    return false;
  }


IsConstantFoldable only applies the “must preserve symbolic contents” guard when GetShapeFromDirectDynamicSource(n, …) is true or the node is explicitly marked _xla_shape_derived. However many shape-tensor plumbing nodes (e.g. StridedSlice, arithmetic ops, etc.) are typically not marked _xla_shape_derived (see MarkForCompilationPass which only marks Shape/ShapeN and Cast). This means these nodes can now become foldable once their Shape input is folded, even in cases where TryGetFoldedValueContents returns false (e.g. an unsupported StridedSlice pattern like negative stride), and the fold will then drop symbolic metadata.

Consider extending the eligibility check so that any node whose (transitive) inputs carry dynamic shape contents is only foldable if TryGetFoldedValueContents succeeds (or alternatively propagate _xla_shape_derived onto these supported plumbing ops during marking/encapsulation).

Copilot · 2026-04-27T11:41:09Z

+    if (has_exact_contents && HasDynamicDimExprs(exact_contents)) {
+      builder.Attr(kUserInferredValueContentsSerializedAttrName,
+                   exact_contents.SerializeAsString());
+    }


The new attrs user_inferred_value_contents_serialized / user_inferred_value_contents are attached to Const nodes, but they are not declared in the Const op definition (see REGISTER_OP("Const") in tensorflow/core/ops/array_ops.cc). When these Const nodes are created/validated (NodeBuilder::Finalize / ValidateNodeDef), TensorFlow will emit periodic ERROR logs for unknown attributes.

To avoid log noise (and make the attr contract explicit), either add these attrs to the Const op registration with appropriate defaults/types, or rename them with a leading _ to mark them as internal optional attrs.

Copilot · 2026-04-27T11:41:09Z

+  if (has_exact_contents && HasDynamicDimExprs(exact_contents)) {
+    builder.Attr("has_dynamic", true)
+        .Attr(kUserInferredValueContentsSerializedAttrName,
+              exact_contents.SerializeAsString());
+  }


Same issue as above: ReplaceTensorWithConstant attaches user_inferred_value_contents_serialized to new Const nodes, but Const's OpDef does not currently declare this attr. This will trigger ValidateNodeDef unknown-attr ERROR logs when these replacement Const nodes are created/validated. Prefer declaring the attr on REGISTER_OP("Const") (or using an internal _... attr name) to avoid operational log spam.

Copilot · 2026-04-27T11:41:09Z

+  const bool has_dynamic =
+      GetShapeFromDirectDynamicSource(n, &exact_contents);


Minor maintainability issue: IsConstantFoldable calls TryGetFoldedValueContents(n, 0, &exact_contents) and then immediately reuses the same exact_contents variable as the output buffer for GetShapeFromDirectDynamicSource(n, &exact_contents). This overwrites the previously-computed contents and is easy to misread (and error-prone if more logic later starts using exact_contents). Prefer using a separate TensorShapeProto dynamic_shape; for the has_dynamic check.

Suggested change

const bool has_dynamic =

GetShapeFromDirectDynamicSource(n, &exact_contents);

TensorShapeProto dynamic_shape;

const bool has_dynamic =

GetShapeFromDirectDynamicSource(n, &dynamic_shape);

Copilot · 2026-04-27T11:41:09Z

+    if (GetNodeAttr(ctx->op_kernel().def(),
+                    "user_inferred_value_contents_serialized",
+                    &inferred_value_contents_serialized)
+            .ok()) {
+      inferred_value_contents_proto.ParseFromString(
+          inferred_value_contents_serialized);
+    } else {
+      GetNodeAttr(ctx->op_kernel().def(), "user_inferred_value_contents",
+                  &inferred_value_contents_proto)
+          .IgnoreError();
+    }


ParseFromString() results are currently ignored for user_inferred_value_contents_serialized. If the attr is present but corrupted/invalid, inferred_value_contents_proto stays empty and the code silently falls back to inferred_shape_proto (or to an empty proto). Consider checking the boolean return of ParseFromString() and treating parse failure as “no contents available” explicitly (e.g., clear the proto / skip attachment) to make the behavior clearer and avoid accidentally acting on partially-parsed data.

Copilot · 2026-04-27T11:41:10Z

+          if (inferred_contents_serialized_it != attr_map.end()) {
+            inferred_shape_proto.ParseFromString(
+                inferred_contents_serialized_it->second.s());
+          } else if (inferred_contents_it != attr_map.end()) {
+            inferred_shape_proto = inferred_contents_it->second.shape();
+          } else {
+            inferred_shape_proto = inferred_shape_it->second.shape();
+          }


Similarly, inferred_shape_proto.ParseFromString(inferred_contents_serialized_it->second.s()) ignores the return value. If parsing fails, inferred_shape_proto will be empty but the code may still proceed for edge cases (e.g., empty-vector constants). Consider checking the parse result and returning early on failure so we only attach expressions when the serialized proto is well-formed.

This reverts commit 56b8330.

This reverts commit ad7ecc2.

stevenvar requested a review from Copilot April 23, 2026 11:26

Copilot started reviewing on behalf of stevenvar April 23, 2026 11:27 View session

Copilot AI reviewed Apr 23, 2026

View reviewed changes

stevenvar force-pushed the steven.partial_symbolic_const_folding branch 4 times, most recently from 74663cc to 3cda438 Compare April 27, 2026 10:49

stevenvar requested a review from Copilot April 27, 2026 11:33

Copilot started reviewing on behalf of stevenvar April 27, 2026 11:34 View session

Copilot AI reviewed Apr 27, 2026

View reviewed changes

stevenvar added 14 commits July 29, 2026 12:21

Add dynamic size expression tests and fix propagation gaps (#27)

c057b93

Support max and conditional dynamic shape expressions (#4)

5e677f7

Preserve dynamic shape expressions across XLA shape reconstruction (#13)

7f3d003

Reject dynamic shape expressions in MLIR XLA kernels (#19)

61d4725

Preserve symbolic contents during partial constant folding

1224069

Serialize symbolic const contents attrs

24c0a4a

Avoid partial non-CPU multi-output constant replacement

4f4f042

Add symbolic constant folding coverage

6035005

Revert "Serialize symbolic const contents attrs"

7e71c2b

This reverts commit 56b8330.

Reapply "Serialize symbolic const contents attrs"

e501ec3

This reverts commit ad7ecc2.

Address symbolic constant folding review feedback

9805707

Simplify symbolic contents attr contract

07bb8e8

Clean up symbolic constant folding

1a38be5

Test symbolic contents through constant folding

0bd94c3

stevenvar force-pushed the steven.partial_symbolic_const_folding branch from b108667 to 0bd94c3 Compare July 29, 2026 13:47

Fix symbolic constant folding tests

129f751

stevenvar force-pushed the steven.partial_symbolic_const_folding branch from e5d3302 to 129f751 Compare July 29, 2026 15:25

stevenvar marked this pull request as draft July 30, 2026 10:53

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow partial symbolic constant folding - #98

Allow partial symbolic constant folding#98
stevenvar wants to merge 15 commits into
joeyye-work:for-serving-2.20from
stevenvar:steven.partial_symbolic_const_folding

stevenvar commented Apr 23, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 27, 2026

Uh oh!

Copilot AI Apr 27, 2026

Uh oh!

Copilot AI Apr 27, 2026

Uh oh!

Copilot AI Apr 27, 2026

Uh oh!

Copilot AI Apr 27, 2026

Uh oh!

Copilot AI Apr 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		const bool has_dynamic =
		GetShapeFromDirectDynamicSource(n, &exact_contents);

Conversation

stevenvar commented Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

stevenvar commented Apr 23, 2026 •

edited

Loading