Preserve dynamic expressions across MLIR XLA boundaries

[stevenvar]

Summary

This PR allows MLIR-only TensorFlow kernels to remain in XLA clusters when their inputs carry dynamic shape expressions.

Before entering the MLIR bridge, TensorFlow shape inference records the logical output expressions. The XLA inputs and argument descriptions are then replaced with expression-free views that keep the same physical dimensions. MLIR therefore lowers an ordinary statically shaped region. When its results return to XLA, the recorded output expressions are restored at the boundary.

Why

The previous behavior rejected every MLIR-only kernel with a dynamic input expression. Passing the expressions through MLIR itself would require broad changes to MLIR and to individual HLO inference rules.

Keeping expressions at the boundary avoids that duplication. It also leaves standard XLA dynamic dimensions disabled and does not introduce SetDimensionSize.

Design notes

• TensorFlow shape inference remains the source of truth for logical output expressions.
• MLIR and the generated internal HLO only see the physical static dimensions.
• Generic MLIR-to-HLO conversion and HLO shape inference are unchanged.
• Inputs marked as compile-time constants keep their existing handling.
• This is a draft because operations whose computation observes a dynamic dimension as a value still need careful validation.

Tests

• Compiles the MLIR-only ReluGrad op with two [8<A>, 4] inputs.
• Checks that the returned XLA shape is still [8<A>, 4].
• Checks that the expression dimension is not converted to standard bounded dynamism.
• Checks that the generated HLO contains no SetDimensionSize instruction.