Propagate symbolic shape contents through TF/XLA

[stevenvar]

• This PR keeps track of values that come from dynamic shapes.
• It passes this information through TensorFlow, tf2xla, and HLO, instead of losing it too early.
• It adds symbolic-content propagation in a few missing places, including Range and scalar binary ops.
• It also stores symbolic contents directly on HLO instructions, so later passes can read them in a structured way.

Turn on with: --tf_xla_enable_symbolic_content=true (off by default)

[Copilot]

Pull request overview

This PR introduces a structured “symbolic contents” metadata channel to preserve and propagate dynamic-shape-derived values through TensorFlow → tf2xla → HLO, enabling later passes to reason about symbolic content without losing it early (guarded by --tf_xla_enable_symbolic_content=true).

Changes:

• Adds contents (repeated ExpressionProto) to HloInstructionProto/HloInstruction and wires it through XlaBuilder serialization/deserialization.
• Replaces the prior string-encoded dynamic-constant frontend-attribute mechanism with structured contents in XLA passes (e.g., DynamicConstantRewriter, HLO constant folding skip logic).
• Propagates symbolic contents across several tf2xla kernels and tf2xla plumbing (XlaExpression, XlaArgument, Range/Shape/Slice/Concat/Pack/binary ops, etc.) plus adds a new JIT flag.

Reviewed changes

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

Show a summary per file

File	Description
third_party/xla/xla/shape_dynexpr.h	Adds missing-expression sentinel and proto→DynExpr conversion helper; improves printing for sentinel.
third_party/xla/xla/shape.h	Uses shared missing-expression sentinel in Shape::expressions.
third_party/xla/xla/shape.cc	Updates DynExpr printing to use print() formatting (vs proto debug).
third_party/xla/xla/service/hlo.proto	Adds contents field (id 92) to HloInstructionProto.
third_party/xla/xla/service/dynamic_constant_rewriter.cc	Rewrites dynamic constants using structured contents instead of frontend attrs.
third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc	Skips folding when constant operands carry dynamic contents.
third_party/xla/xla/hlo/ir/hlo_instruction.h	Stores contents on HloInstruction rare data; exposes getters/setters.
third_party/xla/xla/hlo/ir/hlo_instruction.cc	Serializes/deserializes/prints contents on instructions.
third_party/xla/xla/hlo/builder/xla_builder.h	Adds builder APIs/storage for per-instruction contents.
third_party/xla/xla/hlo/builder/xla_builder.cc	Implements Set/GetInstructionContents and serializes into protos; updates GetDimensionSize handling.
tensorflow/core/framework/tensor_shape_expr.h	Declares DimExpr→DynExpr conversion helper.
tensorflow/core/framework/tensor_shape_expr.cc	Implements DimExpr→DynExpr conversion helper.
tensorflow/core/framework/tensor_shape.h	Uses shared missing-expression sentinel; exposes proto→DynExpr helper decl.
tensorflow/core/framework/tensor_shape.cc	Renames/uses proto→DynExpr helper when reading/writing expressions.
tensorflow/core/framework/BUILD	Adds XLA dep needed by updated tensor shape code.
tensorflow/core/common_runtime/constant_folding.cc	Removes special attrs used to smuggle dynamic shape info into Const nodes.
tensorflow/compiler/tf2xla/xla_op_kernel.cc	Restores per-dim expressions from input “contents” when parsing constant shapes.
tensorflow/compiler/tf2xla/xla_expression.h	Replaces dynamic-constant marker fields with vector-based contents API.
tensorflow/compiler/tf2xla/xla_expression.cc	Implements contents storage via XlaBuilder when enabled, else locally.
tensorflow/compiler/tf2xla/xla_compiler.cc	Preserves symbolic per-element expressions on constant args crossing boundaries.
tensorflow/compiler/tf2xla/xla_argument.h	Replaces single-index dynamic-constant marker with per-element ExpressionProto list.
tensorflow/compiler/tf2xla/xla_argument.cc	Extends equality to include serialized per-element expressions.
tensorflow/compiler/tf2xla/symbolic_content_util.h	Introduces flag check helper + unknown-content sentinel.
tensorflow/compiler/tf2xla/kernels/variable_ops.cc	Attaches symbolic contents to VariableShape outputs.
tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc	Propagates symbolic contents through simple 1D slicing.
tensorflow/compiler/tf2xla/kernels/slice_op.cc	Propagates symbolic contents through simple 1D slice; updates DynamicSlice call.
tensorflow/compiler/tf2xla/kernels/shape_op.cc	Propagates symbolic contents through Shape/ShapeN outputs.
tensorflow/compiler/tf2xla/kernels/sequence_ops.cc	Propagates symbolic contents through Range; supports symbolic size expr for integral iota.
tensorflow/compiler/tf2xla/kernels/pack_op.cc	Propagates symbolic contents through scalar pack/concat patterns.
tensorflow/compiler/tf2xla/kernels/identity_op.cc	Forwards XlaExpression to preserve metadata across Identity.
tensorflow/compiler/tf2xla/kernels/fill_op.cc	Uses symbolic expressions for broadcast dims when available.
tensorflow/compiler/tf2xla/kernels/cwise_ops.h	Adds virtual hook for symbolic propagation in binary ops.
tensorflow/compiler/tf2xla/kernels/cwise_ops.cc	Implements symbolic-content propagation for scalar/1D binary ops when possible.
tensorflow/compiler/tf2xla/kernels/const_op.cc	Removes legacy dynamic-shape Const special-casing.
tensorflow/compiler/tf2xla/kernels/concat_op.cc	Propagates symbolic contents through axis-0 concat for scalar/1D cases.
tensorflow/compiler/tf2xla/kernels/cast_op.cc	Forwards symbolic contents through Cast.
tensorflow/compiler/tf2xla/kernels/binary_ops.cc	Wires symbolic computation for basic arithmetic ops (Add/Sub/Mul/Div/etc.).
tensorflow/compiler/tf2xla/kernels/BUILD	Adds tf2xla common dep for kernels using symbolic-content utilities.
tensorflow/compiler/tf2xla/graph_compiler.cc	Removes legacy dynamic-constant arg plumbing.
tensorflow/compiler/tf2xla/BUILD	Exposes new symbolic-content util header and adds flags dep.
tensorflow/compiler/jit/partially_decluster_pass.cc	Avoids declustering “shape-derived” nodes when symbolic content is enabled.
tensorflow/compiler/jit/mark_for_compilation_pass.cc	Marks Shape/Cast as “shape-derived” only when inputs actually have dynamic expressions.
tensorflow/compiler/jit/kernels/xla_ops.cc	Emits per-element expression protos for rewritten constant args.
tensorflow/compiler/jit/flags.h	Adds tf_xla_enable_symbolic_content flag.
tensorflow/compiler/jit/flags.cc	Registers/parses/initializes the new flag.
tensorflow/compiler/jit/device_compilation_cluster_signature.h	Extends signature to include constant contents in cache keys.
tensorflow/compiler/jit/device_compilation_cluster_signature.cc	Hash/eq/human-string updated for constant contents.

---

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