diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD
index 39f93d17aa2932..a8197744359fde 100644
--- a/tensorflow/compiler/jit/BUILD
+++ b/tensorflow/compiler/jit/BUILD
@@ -780,6 +780,8 @@ cc_library(
         ":flags_headers",
         ":tf_graph_to_hlo_compiler",
         ":xla_compile_util",
+        ":xla_batch_matcher",
+        "//tensorflow/compiler/jit:flags",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/core:framework",
         "//tensorflow/core:framework_lite",
@@ -1006,6 +1008,7 @@ cc_library(
     hdrs = ["shape_inference.h"],
     visibility = [":friends"],
     deps = [
+        ":flags",
         ":shape_inference_helpers",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
@@ -2005,3 +2008,13 @@ tf_cuda_cc_test(
         "@local_xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
     ],
 )
+
+cc_library(
+    name = "xla_batch_matcher",
+    srcs = ["xla_batch_matcher.cc"],
+    hdrs = ["xla_batch_matcher.h"],
+    deps = [
+        "//tensorflow/core/platform:logging",
+        "@local_xla//xla:debug_options_flags",
+    ],
+)
\ No newline at end of file
diff --git a/tensorflow/compiler/jit/device_compilation_cluster_signature.cc b/tensorflow/compiler/jit/device_compilation_cluster_signature.cc
index 81902a28532dbf..5996c5f0b3a70c 100644
--- a/tensorflow/compiler/jit/device_compilation_cluster_signature.cc
+++ b/tensorflow/compiler/jit/device_compilation_cluster_signature.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
 
+#include "absl/strings/str_cat.h"
 #include <string>
 #include <utility>
 #include <variant>
@@ -22,14 +23,23 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 using Signature = DeviceCompilationClusterSignature;
+using ConstantTensor = Signature::ConstantTensor;
 using TensorTypeAndShape = Signature::TensorTypeAndShape;
 
 // Functor that converts a Signature's arg to a human readable string.
 struct SignatureHumanStringAppender {
   explicit SignatureHumanStringAppender(std::string* dest) : dest(dest) {}
   std::string* dest;
-  void operator()(const Tensor& arg) {
-    absl::StrAppend(dest, "; ", arg.DebugString());
+  void operator()(const ConstantTensor& arg) {
+    absl::StrAppend(dest, "; ", arg.value.DebugString());
+    if (!arg.contents.empty()) {
+      absl::StrAppend(dest, " contents=[");
+      for (int i = 0; i < arg.contents.size(); ++i) {
+        if (i > 0) absl::StrAppend(dest, ",");
+        absl::StrAppend(dest, arg.contents[i].DebugString());
+      }
+      absl::StrAppend(dest, "]");
+    }
   }
   void operator()(const TensorTypeAndShape& arg) {
     absl::StrAppend(dest, ",", DataTypeString(arg.first));
@@ -40,18 +50,29 @@ struct SignatureHumanStringAppender {
 // Functor that compares the arg values of two different signatures. Returns
 // true when the args are not equal.
 struct SignatureNotEqual {
-  bool operator()(const Tensor& arg, const Tensor& other) {
-    return arg.dtype() != other.dtype() || arg.shape() != other.shape() ||
-           arg.tensor_data() != other.tensor_data();
+  bool operator()(const ConstantTensor& arg, const ConstantTensor& other) {
+    if (arg.value.dtype() != other.value.dtype() ||
+        arg.value.shape() != other.value.shape() ||
+        arg.value.tensor_data() != other.value.tensor_data() ||
+        arg.contents.size() != other.contents.size()) {
+      return true;
+    }
+    for (int i = 0; i < arg.contents.size(); ++i) {
+      if (arg.contents[i].SerializeAsString() !=
+          other.contents[i].SerializeAsString()) {
+        return true;
+      }
+    }
+    return false;
   }
   bool operator()(const TensorTypeAndShape& arg,
                   const TensorTypeAndShape& other) {
     return arg.first != other.first || arg.second != other.second;
   }
-  bool operator()(const Tensor& arg, const TensorTypeAndShape& other) {
+  bool operator()(const ConstantTensor& arg, const TensorTypeAndShape& other) {
     return true;
   }
-  bool operator()(const TensorTypeAndShape& arg, const Tensor& other) {
+  bool operator()(const TensorTypeAndShape& arg, const ConstantTensor& other) {
     return true;
   }
 };
@@ -61,12 +82,16 @@ struct SignatureNotEqual {
 struct SignatureHashCombiner {
   explicit SignatureHashCombiner(const uint64 h) : h(h) {}
   uint64 h;
-  uint64 operator()(const Tensor& arg) {
-    h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.dtype())));
+  uint64 operator()(const ConstantTensor& arg) {
+    h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.value.dtype())));
     h = Hash64Combine(
-        h, Hash64(arg.tensor_data().data(), arg.tensor_data().size()));
-    for (int dim = 0; dim < arg.dims(); ++dim) {
-      h = Hash64Combine(h, std::hash<int>()(arg.dim_size(dim)));
+        h, Hash64(arg.value.tensor_data().data(), arg.value.tensor_data().size()));
+    for (int dim = 0; dim < arg.value.dims(); ++dim) {
+      h = Hash64Combine(h, std::hash<int>()(arg.value.dim_size(dim)));
+    }
+    for (const xla::ExpressionProto& expr : arg.contents) {
+      std::string serialized = expr.SerializeAsString();
+      h = Hash64Combine(h, Hash64(serialized.data(), serialized.size()));
     }
     return h;
   }
@@ -120,7 +145,8 @@ absl::StatusOr<Signature> Signature::Build(
     switch (arg.kind) {
       case XlaCompiler::Argument::kConstant:
       case XlaCompiler::Argument::kConstantResource:
-        signature.args.push_back(arg.constant_value);
+        signature.args.push_back(
+            ConstantTensor{arg.constant_value, arg.constant_value_expressions});
         break;
       case XlaCompiler::Argument::kParameter:
       case XlaCompiler::Argument::kResource:
diff --git a/tensorflow/compiler/jit/device_compilation_cluster_signature.h b/tensorflow/compiler/jit/device_compilation_cluster_signature.h
index 4acea2a03c2cb4..da2de8de370842 100644
--- a/tensorflow/compiler/jit/device_compilation_cluster_signature.h
+++ b/tensorflow/compiler/jit/device_compilation_cluster_signature.h
@@ -18,8 +18,10 @@ limitations under the License.
 
 #include <utility>
 #include <variant>
+#include <vector>
 
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
 
 namespace tensorflow {
 
@@ -34,7 +36,11 @@ struct DeviceCompilationClusterSignature {
   // argument number. Tensors must be in host memory.
   using TensorTypeAndShape =
       std::pair<DataType, absl::InlinedVector<int64_t, 4>>;
-  absl::InlinedVector<std::variant<Tensor, TensorTypeAndShape>, 8> args;
+  struct ConstantTensor {
+    Tensor value;
+    std::vector<xla::ExpressionProto> contents;
+  };
+  absl::InlinedVector<std::variant<ConstantTensor, TensorTypeAndShape>, 8> args;
 
   bool operator==(const DeviceCompilationClusterSignature& other) const;
 
diff --git a/tensorflow/compiler/jit/device_compilation_profiler.cc b/tensorflow/compiler/jit/device_compilation_profiler.cc
index 5e1b3b26e8ecb5..f8a742ccc01148 100644
--- a/tensorflow/compiler/jit/device_compilation_profiler.cc
+++ b/tensorflow/compiler/jit/device_compilation_profiler.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <utility>
 
 #include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/xla_activity.pb.h"
 #include "tensorflow/compiler/jit/xla_activity_listener.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
@@ -32,9 +33,30 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 bool ShouldBeMegamorphic(int64_t compile_count, int64_t execution_count) {
-  const int64_t kCompileThreshold = 10;
+  int64_t kCompileThreshold = 10;
   const int64_t kMinExecutionsPerCompile = 50;
 
+  int64_t tf_xla_threshold_for_megamorphic =
+      GetMarkForCompilationPassFlags()->tf_xla_threshold_for_megamorphic;
+
+  // Negative values other that -1 cannot be used
+  if (tf_xla_threshold_for_megamorphic < -1) {
+    LOG(FATAL) << "The value for the tf_xla_threshold_for_megamorphic flag "
+               << "is out of range.\n"
+               << "Allowed ranges are (-1) to "
+               << std::numeric_limits<int64_t>::max()
+               << " got " << tf_xla_threshold_for_megamorphic << ".";
+  }
+
+  // -1: setting clusters as Megamorphic is disabled
+  //  0 Default behaviour in Tensorflow
+  //  Any other number sets the compilation threshold
+  if (tf_xla_threshold_for_megamorphic == -1) {
+    return false;
+  } else if (tf_xla_threshold_for_megamorphic > 0) {
+    kCompileThreshold = tf_xla_threshold_for_megamorphic;
+  }
+
   // This heuristic is trying to capture the following property: have we sunk a
   // certain minimum amount of compile time into the cluster that didn't quite
   // "pay off"?
diff --git a/tensorflow/compiler/jit/device_compiler.h b/tensorflow/compiler/jit/device_compiler.h
index 34b22033129b96..a9e36dbc437673 100644
--- a/tensorflow/compiler/jit/device_compiler.h
+++ b/tensorflow/compiler/jit/device_compiler.h
@@ -36,6 +36,7 @@ limitations under the License.
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.h"
 #include "tensorflow/compiler/jit/xla_compile_util.h"
+#include "tensorflow/compiler/jit/xla_batch_matcher.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 #include "tensorflow/core/framework/metrics.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -125,6 +126,7 @@ class DeviceCompiler : public ResourceBase {
   DeviceCompilerClient<ExecutableType, ClientType>* compiler_client() {
     return compiler_client_.get();
   }
+  XlaBatchMatcher* xla_batch_matcher() { return xla_batch_matcher_.get(); }
 
   string DebugString() const override;
 
@@ -177,6 +179,9 @@ class DeviceCompiler : public ResourceBase {
   // Pool of threads for asynchronous compilations.
   std::unique_ptr<thread::ThreadPool> async_compiler_threads_;
 
+  // Specified dynamic batch padding values.
+  std::unique_ptr<XlaBatchMatcher> xla_batch_matcher_;
+
   mutex cluster_mutexes_mu_;
   absl::flat_hash_map<DeviceCompilationClusterSignature, std::unique_ptr<mutex>,
                       DeviceCompilationClusterSignature::Hash>
@@ -225,6 +230,11 @@ DeviceCompiler<ExecutableType, ClientType>::DeviceCompiler(
   async_compiler_threads_ = std::make_unique<tensorflow::thread::ThreadPool>(
       tensorflow::Env::Default(), "async_compiler_threads",
       kNumAsyncDeviceCompilerThreads);
+
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    xla_batch_matcher_ = std::make_unique<XlaBatchMatcher>();
+  }
 }
 
 template <typename ExecutableType, typename ClientType>
diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
index 3e8a43ce08ed58..080bef205441f2 100644
--- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
+++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/optional.h"
+#include "tensorflow/compiler/jit/encapsulate_util.h"
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
@@ -54,9 +55,14 @@ limitations under the License.
 #include "tensorflow/core/public/version.h"
 #include "tensorflow/core/util/device_name_utils.h"
 #include "tensorflow/core/util/dump_graph.h"
-
+#include "tensorflow/core/framework/tensor_shape.pb.h"
 namespace tensorflow {
 
+static const absl::flat_hash_set<absl::string_view> kFailingOps = {
+    "Where",
+    // add more here
+};
+
 const char* const kXlaCompiledKernelAttr = "_XlaCompiledKernel";
 const char* const kXlaNumConstantArgsAttr = "_XlaNumConstantArgs";
 const char* const kXlaNumResourceArgsAttr = "_XlaNumResourceArgs";
@@ -114,6 +120,41 @@ void MarkGuaranteedConstants(
   }
 }
 
+// Helper to convert ExpressionProto to a readable string.
+std::string ExprProtoToString(const ExpressionProto& e) {
+  switch (e.node_type_case()) {
+    case ExpressionProto::kConstantValue:
+      return std::to_string(e.constant_value());
+    case ExpressionProto::kVariableId:
+      return absl::StrCat("Var(", e.variable_id(), ")");
+    case ExpressionProto::kAddNode:
+      return absl::StrCat("(", ExprProtoToString(e.add_node().lhs()), " + ",
+                         ExprProtoToString(e.add_node().rhs()), ")");
+    case ExpressionProto::kSubNode:
+      return absl::StrCat("(", ExprProtoToString(e.sub_node().lhs()), " - ",
+                         ExprProtoToString(e.sub_node().rhs()), ")");
+    case ExpressionProto::kMulNode:
+      return absl::StrCat("(", ExprProtoToString(e.mul_node().lhs()), " * ",
+                         ExprProtoToString(e.mul_node().rhs()), ")");
+    case ExpressionProto::kDivNode:
+      return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ",
+                         ExprProtoToString(e.div_node().rhs()), ")");
+    case ExpressionProto::kMaxNode:
+      return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ",
+                          ExprProtoToString(e.max_node().rhs()), ")");
+    case ExpressionProto::kGtNode:
+      return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ",
+                          ExprProtoToString(e.gt_node().rhs()), ")");
+    case ExpressionProto::kSelectNode:
+      return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()),
+                          ", ", ExprProtoToString(e.select_node().on_true()),
+                          ", ", ExprProtoToString(e.select_node().on_false()),
+                          ")");
+    default:
+      return "<none>";
+  }
+}
+
 struct OutputInputTensorPairHasher {
   uint64 operator()(std::pair<OutputTensor, InputTensor> const& s) const {
     return Hash64Combine(OutputTensor::Hash()(s.first),
@@ -369,6 +410,19 @@ class Encapsulator {
 
 namespace {
 
+bool BuildOutputShapeProto(const Node& node, int output_slot,
+                           TensorShapeProto* proto) {
+  AttrSlice attrs = node.attrs();
+  auto shape_attr =
+      attrs.FindByString(kXlaInferredOutputTensorShapesAttrName);
+  if (shape_attr == nullptr || !shape_attr->has_list() ||
+      shape_attr->list().shape_size() <= output_slot) {
+    return false;
+  }
+  *proto = shape_attr->list().shape(output_slot);
+  return true;
+}
+
 // Return in 'sorted' a topological sort of clusters according to the
 // dependencies encoded in ancestors. clusters is the list of all clusters
 // including clusters that are not present in the ancestors map. has_successors
@@ -470,6 +524,26 @@ absl::Status Encapsulator::Subgraph::RecordArg(
     DataType dtype = edge->dst()->input_type(edge->dst_input());
     builder.Attr("T", dtype);
     builder.Attr("index", arg_index);
+    AttrSlice attrs = src_node->attrs();
+    TensorShapeProto output_shape_proto;
+    if (BuildOutputShapeProto(*src_node, src_slot, &output_shape_proto)) {
+      VLOG(1) << "Adding following output shapes for node " << src_node->name()
+              << " : " << output_shape_proto.DebugString();
+      builder.Attr("_output_shapes", {output_shape_proto});
+      builder.Attr(kXlaInferredOutputShapesAttrName, {output_shape_proto});
+    } else {
+      // if cluster argument is the real argument.
+      auto build_attr = attrs.FindByString("_dynamic_dim");
+      if (build_attr) {
+        VLOG(1) << "Found Dynamic dimension in " << src_node->name() << ":"
+                << src_slot;
+        builder.Attr("_dynamic_dim", *build_attr);
+      }
+    }
+    auto shape_derived_attr = attrs.FindByString(kXlaShapeDerivedAttrName);
+    if (shape_derived_attr) {
+      builder.Attr(kXlaShapeDerivedAttrName, *shape_derived_attr);
+    }
     absl::Status s = builder.Finalize(&arg_def);
     if (!s.ok()) return s;
 
@@ -1143,6 +1217,14 @@ static absl::Status RenumberArguments(Graph* graph,
   return absl::OkStatus();
 }
 
+static bool SubgraphHasFailingOps(const Graph& g) {
+  for (Node* n : g.op_nodes()) {
+    if (n->IsRetval()) continue;
+    if (kFailingOps.contains(n->def().op())) return true;
+  }
+  return false;
+}
+
 absl::Status EncapsulateSubgraphsPass::Run(
     const GraphOptimizationPassOptions& options) {
   VLOG(1) << "EncapsulateSubgraphsPass::Run";
@@ -1289,8 +1371,8 @@ absl::Status EncapsulateSubgraphsPass::Run(
 
         // TODO(phawkins): add a forward is-constant analysis, similarly split
         // outputs into host-memory constants and device-memory non-constants.
-
-        AddNodeAttr(kXlaCompiledKernelAttr, true, node);
+        bool compile_enabled = !SubgraphHasFailingOps(**subgraph);
+        AddNodeAttr(kXlaCompiledKernelAttr, compile_enabled, node);
         AddNodeAttr(kXlaNumConstantArgsAttr, num_consts, node);
         AddNodeAttr(kXlaNumResourceArgsAttr, num_resources, node);
         return absl::OkStatus();
diff --git a/tensorflow/compiler/jit/encapsulate_util.cc b/tensorflow/compiler/jit/encapsulate_util.cc
index fa94a341bbabc6..e680849a0fb8f6 100644
--- a/tensorflow/compiler/jit/encapsulate_util.cc
+++ b/tensorflow/compiler/jit/encapsulate_util.cc
@@ -303,6 +303,11 @@ absl::Status PostprocessControlEdgesBetweenOutsideCompilations(
 }  // namespace
 
 const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes";
+const char kXlaInferredOutputTensorShapesAttrName[] =
+    "_xla_inferred_output_tensor_shapes";
+const char kXlaInferredOutputShapesAttrName[] =
+    "_xla_inferred_output_shapes";
+const char kXlaShapeDerivedAttrName[] = "_xla_shape_derived";
 
 const char kXlaConnectedToXlaComputationAttrName[] =
     "_xla_connected_to_xla_computation";
diff --git a/tensorflow/compiler/jit/encapsulate_util.h b/tensorflow/compiler/jit/encapsulate_util.h
index 7c99763c770728..f0f0b0570fd19c 100644
--- a/tensorflow/compiler/jit/encapsulate_util.h
+++ b/tensorflow/compiler/jit/encapsulate_util.h
@@ -28,6 +28,19 @@ namespace tensorflow {
 // a list of PartialTensorShape objects.
 extern const char kXlaInferredShapesAttrName[];
 
+// Attribute marking Grappler-inferred output TensorShapeProtos. Attribute
+// value is a list of TensorShapeProto objects and may include ExpressionProto
+// annotations when available.
+extern const char kXlaInferredOutputTensorShapesAttrName[];
+
+// Attribute carrying inferred output TensorShapeProtos on encapsulated _Arg
+// nodes for XlaCompileOp argument reconstruction.
+extern const char kXlaInferredOutputShapesAttrName[];
+
+// Attribute indicating that a node is shape-derived and worth tracking across
+// JIT boundaries, e.g. Shape and shape-fed Cast.
+extern const char kXlaShapeDerivedAttrName[];
+
 // Infers output shapes for all nodes in graph `g`. The output shapes will be
 // stored in node attribute `kXlaInferredShapesAttrName`.
 //
diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc
index 10756ddf9de7b5..5a6f741a01e972 100644
--- a/tensorflow/compiler/jit/flags.cc
+++ b/tensorflow/compiler/jit/flags.cc
@@ -108,6 +108,15 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector<Flag>* flag_list) {
       Flag("tf_xla_max_cluster_size",
            &mark_for_compilation_flags->tf_xla_max_cluster_size,
            "Maximum number of operators in an XLA compilation."),
+      Flag("tf_xla_annotate_cluster_id",
+           &mark_for_compilation_flags->tf_xla_annotate_cluster_id,
+           "Allow operator names to influence clustering scheme."
+           "Operators whose name starting with .cluster.{id} will likely"
+           "to be clustered together if the ids are the same number. "
+           ".cluster.none will not be clustered with those having numbered id"),
+      Flag("tf_xla_cluster_parallel",
+           &mark_for_compilation_flags->tf_xla_cluster_parallel,
+           "Split parallel compute subgraph info different clusters"),
       Flag(
           "tf_xla_ops_to_cluster",
           &mark_for_compilation_flags->tf_xla_ops_to_cluster,
@@ -155,6 +164,12 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector<Flag>* flag_list) {
            &mark_for_compilation_flags->tf_xla_deterministic_cluster_names,
            "Causes the function names assigned by auto clustering to be "
            "deterministic from run to run."),
+      Flag("tf_xla_enable_dynamic_sizes",
+           &mark_for_compilation_flags->tf_xla_enable_dynamic_sizes,
+           "Enable dynamic sizes support."),
+      Flag("tf_xla_enable_symbolic_content",
+           &mark_for_compilation_flags->tf_xla_enable_symbolic_content,
+           "Enable symbolic content propagation."),
       Flag("tf_xla_persistent_cache_directory",
            &mark_for_compilation_flags->tf_xla_persistent_cache_directory,
            "If non-empty, JIT-compiled executables are saved to and loaded "
@@ -175,6 +190,11 @@ void AppendMarkForCompilationPassFlagsInternal(std::vector<Flag>* flag_list) {
            &mark_for_compilation_flags->tf_xla_persistent_cache_prefix,
            "Specifies the persistance cache prefix. Default is "
            "\"xla_compile_cache\""),
+     Flag("tf_xla_threshold_for_megamorphic",
+           &mark_for_compilation_flags->tf_xla_threshold_for_megamorphic,
+           "Sets the threshold for marking a cluster megamorphic. "
+           "Setting it to -1 disables marking clusters megamorphic."
+           "Setting it to 0 uses the default behaviour of TensorFlow."),
       Flag("tf_xla_sparse_core_disable_table_stacking",
            &sparse_core_flags->tf_xla_sparse_core_disable_table_stacking,
            "Disable table stacking for all the tables passed to the SparseCore"
@@ -232,15 +252,20 @@ void AllocateAndParseFlags() {
   mark_for_compilation_flags->tf_xla_min_cluster_size = 4;
   mark_for_compilation_flags->tf_xla_max_cluster_size =
       std::numeric_limits<int32>::max();
+  mark_for_compilation_flags->tf_xla_annotate_cluster_id = false;
+  mark_for_compilation_flags->tf_xla_cluster_parallel = false;
   mark_for_compilation_flags->tf_xla_clustering_debug = false;
   mark_for_compilation_flags->tf_xla_cpu_global_jit = false;
   mark_for_compilation_flags->tf_xla_clustering_fuel =
       std::numeric_limits<int64_t>::max();
+  mark_for_compilation_flags->tf_xla_threshold_for_megamorphic = 0;
   mark_for_compilation_flags
       ->tf_xla_disable_deadness_safety_checks_for_debugging = false;
   mark_for_compilation_flags
       ->tf_xla_disable_resource_variable_safety_checks_for_debugging = false;
   mark_for_compilation_flags->tf_xla_deterministic_cluster_names = false;
+  mark_for_compilation_flags->tf_xla_enable_dynamic_sizes = false;
+  mark_for_compilation_flags->tf_xla_enable_symbolic_content = false;
   mark_for_compilation_flags->tf_xla_persistent_cache_directory = "";
   mark_for_compilation_flags->tf_xla_persistent_cache_device_types = "";
   mark_for_compilation_flags->tf_xla_persistent_cache_read_only = false;
@@ -404,7 +429,8 @@ void AllocateAndParseFlags() {
             "and creates TPUReshardVariables ops.")});
 
   AppendMarkForCompilationPassFlagsInternal(flag_list);
-  xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list);
+  xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list,
+                                        /*reset_envvar=*/true);
 
   mlir_flags = new MlirCommonFlags;
   if (!enable_mlir_bridge_is_explicit) {
diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h
index 0d0c5082cf9a82..480e20dc7ec296 100644
--- a/tensorflow/compiler/jit/flags.h
+++ b/tensorflow/compiler/jit/flags.h
@@ -62,6 +62,12 @@ struct MarkForCompilationPassFlags {
   // Maximum number of operators in an XLA compilation.
   int32 tf_xla_max_cluster_size;
 
+  // Enable operator name to influence clustering decision
+  bool tf_xla_annotate_cluster_id;
+
+  // Split parallel compute subgraph info different clusters
+  bool tf_xla_cluster_parallel;
+
   // If non-empty, limit XLA clustering to the following TF operations.
   string tf_xla_ops_to_cluster;
 
@@ -93,6 +99,12 @@ struct MarkForCompilationPassFlags {
   // so that they remain stable from run to run of auto clusteing.
   bool tf_xla_deterministic_cluster_names;
 
+  // If true enables support of dynamic sizes.
+  bool tf_xla_enable_dynamic_sizes;
+
+  // If true enables symbolic content propagation.
+  bool tf_xla_enable_symbolic_content;
+
   // If non-empty, JIT-compiled executables are saved to and loaded from the
   // specified file system directory path.
   std::string tf_xla_persistent_cache_directory;
@@ -111,6 +123,11 @@ struct MarkForCompilationPassFlags {
 
   // Specifies the persistance cache prefix. Default is "xla_compile_cache"
   string tf_xla_persistent_cache_prefix;
+
+  // Sets the threshold for marking a cluster megamorphic.
+  // Setting it to -1 disables marking clusters megamorphic.
+  // Setting it to 0 uses the default behaviour of TensorFlow.
+  int64_t tf_xla_threshold_for_megamorphic;
 };
 
 // Flags associated with XLA Sparse Core.
diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD
index dec057ebfba9ab..1af416550a43dd 100644
--- a/tensorflow/compiler/jit/kernels/BUILD
+++ b/tensorflow/compiler/jit/kernels/BUILD
@@ -1,3 +1,4 @@
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 
 package(
@@ -48,7 +49,10 @@ XLA_OPS_DEPS = [
 cc_library(
     name = "xla_ops_no_jit_rewrite_registration",
     srcs = ["xla_ops.cc"],
-    hdrs = ["xla_ops.h"],
+    hdrs = [
+        "xla_ops.h",
+        "xla_ops_internal.h",
+    ],
     deps = XLA_OPS_DEPS + [
         "//tensorflow/compiler/jit:device_compilation_cache",
         "//tensorflow/compiler/jit:device_compilation_profiler",
@@ -56,10 +60,12 @@ cc_library(
         "//tensorflow/compiler/jit:tf_graph_to_hlo_compiler",
         "//tensorflow/compiler/jit:tf_to_hlo_compiler",
         "//tensorflow/compiler/jit:xla_compile_util",
+        "//tensorflow/compiler/jit:xla_batch_matcher",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/platform:refcount",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:node_hash_map",
+        "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -74,6 +80,20 @@ cc_library(
     alwayslink = 1,
 )
 
+tf_cc_test(
+    name = "xla_ops_test",
+    srcs = ["xla_ops_test.cc"],
+    env = {"TF_XLA_FLAGS": "--tf_xla_enable_dynamic_sizes"},
+    deps = [
+        ":xla_ops_no_jit_rewrite_registration",
+        "//tensorflow/compiler/tf2xla:xla_argument",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/platform:test",
+        "@com_google_googletest//:gtest_main",
+        "@local_xla//xla:shape_util",
+    ],
+)
+
 cc_library(
     name = "xla_ops",
     hdrs = ["xla_ops.h"],
diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc
index 468b85280e2a47..c28c8c3400414a 100644
--- a/tensorflow/compiler/jit/kernels/xla_ops.cc
+++ b/tensorflow/compiler/jit/kernels/xla_ops.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/compiler/jit/kernels/xla_ops.h"
+#include "tensorflow/compiler/jit/kernels/xla_ops_internal.h"
 
 #include <cstdint>
 #include <functional>
@@ -36,12 +37,14 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "tensorflow/compiler/jit/device_compilation_profiler.h"
 #include "tensorflow/compiler/jit/device_compiler.h"
 #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h"
+#include "tensorflow/compiler/jit/encapsulate_util.h"
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/pjrt_compile_util.h"
 #include "tensorflow/compiler/jit/variable_info.h"
@@ -54,6 +57,7 @@ limitations under the License.
 #include "tensorflow/compiler/jit/xla_host_send_device_context.h"
 #include "tensorflow/compiler/jit/xla_launch_util.h"
 #include "tensorflow/compiler/jit/xla_platform_info.h"
+#include "tensorflow/compiler/jit/xla_batch_matcher.h"
 #include "tensorflow/compiler/tf2xla/tf2xla_util.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
@@ -61,14 +65,18 @@ limitations under the License.
 #include "xla/client/local_client.h"
 #include "xla/executable_run_options.h"
 #include "xla/pjrt/pjrt_client.h"
+#include "xla/printer.h"
 #include "xla/service/gpu/gpu_executable_run_options.h"
+#include "xla/shape_expr.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/protobuf/error_codes.pb.h"
 #include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/batch_size_resource.h"
 #include "tensorflow/core/framework/node_def_util.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/monitoring/counter.h"
@@ -95,6 +103,8 @@ limitations under the License.
 namespace tensorflow {
 
 namespace {
+constexpr char kUserInferredValueContentsAttrName[] =
+    "_user_inferred_value_contents";
 using XlaDeviceCompiler =
     DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
 using PjRtDeviceCompiler =
@@ -199,6 +209,571 @@ using PjRtExecutableClosure =
 using PjRtExecutableClosureStore =
     ExecutableClosureStore<xla::PjRtLoadedExecutable, xla::PjRtClient>;
 
+struct DynamicSolveEvidence {
+  int64_t solved_value;
+  int64_t observed_value;
+  std::string expr;
+};
+
+struct DynamicSolveSelection {
+  std::set<int64_t> singleton_values;
+  std::set<int64_t> nonsingleton_values;
+  std::set<std::string> expressions;
+  std::optional<int64_t> chosen_value;
+  std::string reason;
+  bool rejected = false;
+};
+
+template <typename Candidate>
+DynamicSolveSelection SelectDynamicSolveEvidence(
+    absl::Span<const Candidate> candidates) {
+  DynamicSolveSelection selection;
+  for (const Candidate& candidate : candidates) {
+    const DynamicSolveEvidence& evidence = candidate.evidence;
+    selection.expressions.insert(evidence.expr);
+    if (evidence.observed_value == 1) {
+      selection.singleton_values.insert(evidence.solved_value);
+    } else {
+      selection.nonsingleton_values.insert(evidence.solved_value);
+    }
+  }
+
+  // A runtime singleton may be an implicitly broadcast operand. Prefer
+  // consistent non-singleton evidence for the same normalized variable.
+  if (!selection.nonsingleton_values.empty()) {
+    if (selection.nonsingleton_values.size() != 1) {
+      selection.rejected = true;
+      selection.reason = "conflicting non-singleton candidates";
+      return selection;
+    }
+    selection.chosen_value = *selection.nonsingleton_values.begin();
+    selection.reason =
+        selection.singleton_values.empty()
+            ? "used non-singleton evidence"
+            : "preferred non-singleton evidence over singleton-derived "
+              "candidates";
+    return selection;
+  }
+
+  if (selection.singleton_values.size() == 1) {
+    selection.chosen_value = *selection.singleton_values.begin();
+    selection.reason = "used singleton-derived evidence only";
+  } else if (selection.singleton_values.size() > 1) {
+    selection.rejected = true;
+    selection.reason = "conflicting singleton-derived candidates";
+  } else {
+    selection.reason = "no dynamic values were solved";
+  }
+  return selection;
+}
+
+struct DynamicSolveCandidate {
+  DynamicSolveEvidence evidence;
+  std::string context;
+};
+
+}  // namespace
+
+namespace xla_ops_internal {
+
+int GetRuntimeInputIndex(absl::Span<const int> input_mapping,
+                         int xla_input_index, int num_constant_args,
+                         bool constants_omitted) {
+  if (xla_input_index < 0 || xla_input_index >= input_mapping.size()) {
+    return -1;
+  }
+  const int missing_input_prefix = constants_omitted ? num_constant_args : 0;
+  return input_mapping[xla_input_index] - missing_input_prefix;
+}
+
+std::optional<int64_t> GetConstantArgumentElementValue(
+    const XlaArgument& arg, int index) {
+  if (index < 0 || index >= arg.constant_value.NumElements()) {
+    return std::nullopt;
+  }
+  switch (arg.constant_value.dtype()) {
+    case DT_INT32:
+      return static_cast<int64_t>(arg.constant_value.flat<int32>()(index));
+    case DT_INT64:
+      return static_cast<int64_t>(arg.constant_value.flat<int64_t>()(index));
+    default:
+      return std::nullopt;
+  }
+}
+
+void SetConstantArgumentExpressionToLiteralValue(XlaArgument* arg,
+                                                 int index) {
+  if (arg == nullptr || index < 0 ||
+      index >= arg->constant_value_expressions.size()) {
+    return;
+  }
+  std::optional<int64_t> value = GetConstantArgumentElementValue(*arg, index);
+  if (!value.has_value()) {
+    arg->constant_value_expressions[index].Clear();
+    return;
+  }
+  arg->constant_value_expressions[index].Clear();
+  arg->constant_value_expressions[index].set_constant_value(*value);
+}
+
+DynamicSolveFilterDecision AnalyzeIgnoredDynamicArgumentOccurrences(
+    absl::Span<const XlaArgument> args) {
+  struct Candidate {
+    DynamicSolveEvidence evidence;
+    IgnoredDynamicArgumentOccurrence occurrence;
+  };
+
+  DynamicSolveFilterDecision result;
+  std::map<std::set<int>, std::vector<Candidate>> variable_ids_to_candidates;
+
+  for (int arg_index = 0; arg_index < args.size(); ++arg_index) {
+    const XlaArgument& arg = args[arg_index];
+    if (absl::holds_alternative<TensorShape>(arg.shape)) {
+      const TensorShape& shape = std::get<TensorShape>(arg.shape);
+      for (int dim = 0; dim < shape.get_expressions().size(); ++dim) {
+        const xla::DExpr& expr = shape.get_expression(dim);
+        if (!(expr && expr->is_dynamic())) {
+          continue;
+        }
+        xla::DExpr simplified_expr = expr.simplify();
+        std::optional<int64_t> solved_value =
+            simplified_expr->solve(shape.dim_size(dim));
+        if (!solved_value.has_value()) {
+          continue;
+        }
+        const std::string expr_string = DExprToString(simplified_expr);
+        variable_ids_to_candidates[simplified_expr->get_all_ids()].push_back(
+            Candidate{
+                DynamicSolveEvidence{*solved_value, shape.dim_size(dim),
+                                     expr_string},
+                IgnoredDynamicArgumentOccurrence{
+                    IgnoredDynamicArgumentOccurrence::Source::kShapeDimension,
+                    arg_index,
+                    dim,
+                    shape.dim_size(dim),
+                    *solved_value,
+                    expr_string}});
+      }
+    }
+
+    if (arg.kind != XlaCompiler::Argument::kConstant ||
+        arg.constant_value_expressions.empty()) {
+      continue;
+    }
+    for (int element_index = 0;
+         element_index < arg.constant_value_expressions.size();
+         ++element_index) {
+      xla::DExpr expr =
+          xla::DExprFromProto(arg.constant_value_expressions[element_index]);
+      if (!(expr && expr->is_dynamic())) {
+        continue;
+      }
+      std::optional<int64_t> observed_value =
+          GetConstantArgumentElementValue(arg, element_index);
+      if (!observed_value.has_value()) {
+        continue;
+      }
+      xla::DExpr simplified_expr = expr.simplify();
+      std::optional<int64_t> solved_value =
+          simplified_expr->solve(*observed_value);
+      if (!solved_value.has_value()) {
+        continue;
+      }
+      const std::string expr_string = DExprToString(simplified_expr);
+      variable_ids_to_candidates[simplified_expr->get_all_ids()].push_back(
+          Candidate{
+              DynamicSolveEvidence{*solved_value, *observed_value,
+                                   expr_string},
+              IgnoredDynamicArgumentOccurrence{
+                  IgnoredDynamicArgumentOccurrence::Source::
+                      kConstantValueElement,
+                  arg_index,
+                  element_index,
+                  *observed_value,
+                  *solved_value,
+                  expr_string}});
+    }
+  }
+
+  std::vector<std::string> expr_summaries;
+  for (const auto& [variable_ids, candidates] : variable_ids_to_candidates) {
+    DynamicSolveSelection selection = SelectDynamicSolveEvidence(
+        absl::MakeConstSpan(candidates));
+
+    if (selection.rejected) {
+      result.can_run = false;
+    } else if (selection.chosen_value.has_value() &&
+               !selection.nonsingleton_values.empty()) {
+      for (const auto& candidate : candidates) {
+        if (candidate.evidence.observed_value == 1 &&
+            candidate.evidence.solved_value != *selection.chosen_value) {
+          VLOG(1) << "Ignoring singleton-derived dynamic "
+                  << "occurrence during XLA signature filtering: "
+                  << "variable_ids={" << absl::StrJoin(variable_ids, ", ")
+                  << "} expr=" << candidate.occurrence.expr
+                  << " chosen_value=" << *selection.chosen_value
+                  << " ignored_observed_value="
+                  << candidate.evidence.observed_value
+                  << " ignored_solved_value="
+                  << candidate.evidence.solved_value
+                  << " arg_index=" << candidate.occurrence.arg_index
+                  << (candidate.occurrence.source ==
+                              IgnoredDynamicArgumentOccurrence::Source::
+                                  kShapeDimension
+                          ? " dim="
+                          : " element=")
+                  << candidate.occurrence.dim_or_index;
+          result.ignored_occurrences.push_back(candidate.occurrence);
+        }
+      }
+    }
+
+    expr_summaries.push_back(absl::StrCat(
+        "variable_ids={", absl::StrJoin(variable_ids, ", "),
+        "} expressions={", absl::StrJoin(selection.expressions, ", "),
+        "} chosen=",
+        selection.chosen_value.has_value()
+            ? std::to_string(*selection.chosen_value)
+            : std::string("<none>"),
+        " singleton_derived_values={",
+        absl::StrJoin(selection.singleton_values, ", "),
+        "} nonsingleton_values={",
+        absl::StrJoin(selection.nonsingleton_values, ", "), "} reason=",
+        selection.reason));
+  }
+
+  if (!result.can_run) {
+    result.diagnostic = absl::StrCat(
+        "Failed to recover a unique XLA dynamic batch size after solve-time "
+        "candidate filtering. solved_expressions=[",
+        absl::StrJoin(expr_summaries, " | "), "]");
+    return result;
+  }
+
+  if (!result.ignored_occurrences.empty()) {
+    std::vector<std::string> ignored_summaries;
+    ignored_summaries.reserve(result.ignored_occurrences.size());
+    for (const auto& occurrence : result.ignored_occurrences) {
+      ignored_summaries.push_back(absl::StrCat(
+          occurrence.source ==
+                  IgnoredDynamicArgumentOccurrence::Source::kShapeDimension
+              ? "shape_arg="
+              : "const_arg=",
+          occurrence.arg_index,
+          occurrence.source ==
+                  IgnoredDynamicArgumentOccurrence::Source::kShapeDimension
+              ? " dim="
+              : " element=",
+          occurrence.dim_or_index, " expr=", occurrence.expr,
+          " observed=", occurrence.observed_value,
+          " solved=", occurrence.solved_value));
+    }
+    result.diagnostic = absl::StrCat(
+        "Ignoring singleton-derived dynamic occurrences for XLA "
+        "signature/HLO: ",
+        absl::StrJoin(ignored_summaries, "; "));
+  }
+  return result;
+}
+
+std::vector<XlaArgument> BuildStaticCompilationArguments(
+    absl::Span<const XlaArgument> args) {
+  std::vector<XlaArgument> static_args(args.begin(), args.end());
+  auto clear_shape_expressions = [&](auto&& self, xla::Shape* shape) -> void {
+    if (shape->IsTuple()) {
+      for (xla::Shape& subshape : *shape->mutable_tuple_shapes()) {
+        self(self, &subshape);
+      }
+    } else if (shape->IsArray()) {
+      shape->set_expressions({});
+      for (int dim = 0; dim < shape->dimensions().size(); ++dim) {
+        if (shape->dimensions(dim) >= 0) {
+          shape->set_dynamic_dimension(dim, false);
+        }
+      }
+    }
+  };
+
+  // Preserve concrete runtime dimensions and values so CompileIfNeeded uses
+  // its ordinary static-shape cache key. Only symbolic/dynamic annotations
+  // are removed.
+  for (XlaArgument& arg : static_args) {
+    if (absl::holds_alternative<TensorShape>(arg.shape)) {
+      std::get<TensorShape>(arg.shape).set_expressions({});
+    } else {
+      clear_shape_expressions(clear_shape_expressions,
+                              &std::get<xla::Shape>(arg.shape));
+    }
+    arg.constant_value_expressions.clear();
+  }
+  return static_args;
+}
+
+void StripIgnoredDynamicArgumentOccurrences(
+    const std::vector<IgnoredDynamicArgumentOccurrence>& ignored_occurrences,
+    std::vector<XlaArgument>* args) {
+  if (args == nullptr) {
+    return;
+  }
+  for (const auto& occurrence : ignored_occurrences) {
+    if (occurrence.arg_index < 0 || occurrence.arg_index >= args->size()) {
+      continue;
+    }
+    XlaArgument& arg = (*args)[occurrence.arg_index];
+    if (occurrence.source ==
+        IgnoredDynamicArgumentOccurrence::Source::kShapeDimension) {
+      if (!absl::holds_alternative<TensorShape>(arg.shape)) {
+        continue;
+      }
+      VLOG(1) << "Dropping ignored dynamic shape expression from XLA "
+              << "signature/HLO: arg_index=" << occurrence.arg_index
+              << " dim=" << occurrence.dim_or_index
+              << " expr=" << occurrence.expr
+              << " observed_value=" << occurrence.observed_value
+              << " solved_value=" << occurrence.solved_value;
+      TensorShape& shape = std::get<TensorShape>(arg.shape);
+      shape.set_expression(occurrence.dim_or_index, xla::DExpr());
+      continue;
+    }
+
+    VLOG(1) << "Dropping ignored dynamic constant-value expression from XLA "
+            << "signature/HLO: arg_index=" << occurrence.arg_index
+            << " element=" << occurrence.dim_or_index
+            << " expr=" << occurrence.expr
+            << " observed_value=" << occurrence.observed_value
+            << " solved_value=" << occurrence.solved_value;
+    SetConstantArgumentExpressionToLiteralValue(&arg, occurrence.dim_or_index);
+  }
+}
+
+DynamicBatchResolutionResult ResolveDynamicBatchSizeFromRuntimeShapes(
+    absl::Span<const xla::Shape> xla_input_shapes,
+    int runtime_input_count,
+    absl::FunctionRef<int(int)> get_runtime_input_index,
+    absl::FunctionRef<const TensorShape&(int)> get_runtime_input_shape,
+    absl::FunctionRef<absl::string_view(int)> get_runtime_input_name,
+    bool log_solves, absl::string_view cluster_name) {
+  DynamicBatchResolutionResult result;
+  std::set<int64_t> dyn_vals;
+  std::map<std::set<int>, std::vector<DynamicSolveCandidate>>
+      variable_ids_to_candidates;
+
+  for (int i = 0; i < xla_input_shapes.size(); ++i) {
+    const xla::Shape& xla_shape = xla_input_shapes[i];
+    const int input_idx = get_runtime_input_index(i);
+    if (!xla_shape.IsArray() || xla_shape.expressions().empty() ||
+        input_idx < 0 || input_idx >= runtime_input_count) {
+      continue;
+    }
+
+    const TensorShape& runtime_shape = get_runtime_input_shape(input_idx);
+    const absl::string_view input_name = get_runtime_input_name(input_idx);
+    for (int dim = 0; dim < xla_shape.expressions().size(); ++dim) {
+      const xla::DExpr& expr = xla_shape.expressions(dim);
+      if (!(expr && expr->is_dynamic())) {
+        continue;
+      }
+      xla::DExpr simplified_expr = expr.simplify();
+      const int64_t size = runtime_shape.dim_size(dim);
+      std::optional<int64_t> dyn_val = simplified_expr->solve(size);
+      if (log_solves && dyn_val.has_value()) {
+        VLOG(1) << "Dynamic solve: cluster=" << cluster_name
+                << " runtime_input_index=" << input_idx
+                << " xla_input_index=" << i << " dim=" << dim
+                << " input_name=" << input_name
+                << " expr=" << DExprToString(simplified_expr)
+                << " target_size=" << size << " result=" << *dyn_val;
+      }
+      if (!dyn_val.has_value()) {
+        if (log_solves) {
+          LOG(WARNING) << "Dynamic solve failed: cluster=" << cluster_name
+                       << " runtime_input_index=" << input_idx
+                       << " xla_input_index=" << i << " dim=" << dim
+                       << " input_name=" << input_name
+                       << " expr=" << DExprToString(simplified_expr)
+                       << " target_size=" << size;
+        }
+        continue;
+      }
+      const std::string expr_string = DExprToString(simplified_expr);
+      const std::string context = absl::StrCat(
+          "xla_input_index=", i, " runtime_input_index=", input_idx,
+          " input_name=", input_name, " dim=", dim,
+          " runtime_dim_size=", size,
+          " solved_dynamic_value=", *dyn_val);
+      variable_ids_to_candidates[simplified_expr->get_all_ids()].push_back(
+          DynamicSolveCandidate{
+              DynamicSolveEvidence{*dyn_val, size, expr_string}, context});
+    }
+  }
+
+  std::vector<std::string> expr_summaries;
+  std::optional<std::set<int>> expected_dyn_ids;
+  bool mismatched_dyn_ids = false;
+  bool candidate_filter_rejected = false;
+  for (const auto& [variable_ids, candidates] : variable_ids_to_candidates) {
+    std::vector<std::string> singleton_contexts;
+    std::vector<std::string> nonsingleton_contexts;
+    for (const auto& candidate : candidates) {
+      if (candidate.evidence.observed_value == 1) {
+        singleton_contexts.push_back(candidate.context);
+      } else {
+        nonsingleton_contexts.push_back(candidate.context);
+      }
+    }
+    DynamicSolveSelection selection =
+        SelectDynamicSolveEvidence(absl::MakeConstSpan(candidates));
+
+    if (selection.rejected) {
+      candidate_filter_rejected = true;
+      LOG(WARNING) << "Dynamic solve rejected candidates for variable_ids={"
+                   << absl::StrJoin(variable_ids, ", ")
+                   << "}: " << selection.reason;
+    }
+    if (selection.chosen_value.has_value()) {
+      dyn_vals.insert(*selection.chosen_value);
+    }
+
+    expr_summaries.push_back(absl::StrCat(
+        "variable_ids={", absl::StrJoin(variable_ids, ", "),
+        "} expressions={", absl::StrJoin(selection.expressions, ", "),
+        "} chosen=",
+        selection.chosen_value.has_value()
+            ? std::to_string(*selection.chosen_value)
+            : std::string("<none>"),
+        " singleton_derived_values={",
+        absl::StrJoin(selection.singleton_values, ", "),
+        "} nonsingleton_values={",
+        absl::StrJoin(selection.nonsingleton_values, ", "),
+        "} reason=", selection.reason, " singleton_derived_contexts=[",
+        absl::StrJoin(singleton_contexts, "; "),
+        "] nonsingleton_contexts=[",
+        absl::StrJoin(nonsingleton_contexts, "; "), "]"));
+  }
+
+  for (const xla::Shape& xla_shape : xla_input_shapes) {
+    if (!xla_shape.IsArray() || xla_shape.expressions().empty()) {
+      continue;
+    }
+    for (const xla::DExpr& expr : xla_shape.expressions()) {
+      if (!(expr && expr->is_dynamic())) {
+        continue;
+      }
+      std::set<int> ids = expr->get_all_ids();
+      if (!expected_dyn_ids.has_value()) {
+        expected_dyn_ids = ids;
+      } else if (*expected_dyn_ids != ids) {
+        mismatched_dyn_ids = true;
+      }
+    }
+  }
+
+  if (dyn_vals.size() == 1 && expected_dyn_ids.has_value() &&
+      expected_dyn_ids->size() == 1 && !mismatched_dyn_ids) {
+    result.has_batch_size = true;
+    result.batch_size = *dyn_vals.begin();
+    result.diagnostic = absl::StrCat(
+        "shared variable ids={", absl::StrJoin(*expected_dyn_ids, ", "),
+        "} value=", result.batch_size);
+    return result;
+  }
+
+  if (variable_ids_to_candidates.empty()) {
+    result.diagnostic = "no dynamic values were solved";
+  } else if (candidate_filter_rejected) {
+    result.can_run = false;
+    result.diagnostic = absl::StrCat(
+        "Failed to recover a unique XLA dynamic batch size after solve-time "
+        "candidate filtering. solved_expressions=[",
+        absl::StrJoin(expr_summaries, " | "), "] all_values={",
+        absl::StrJoin(dyn_vals, ", "), "}");
+  } else if (dyn_vals.size() == 1 && mismatched_dyn_ids) {
+    result.diagnostic = absl::StrCat(
+        "solved dynamic expressions do not share the same variable ids: ",
+        absl::StrJoin(expr_summaries, " | "));
+  } else if (dyn_vals.size() == 1) {
+    result.diagnostic = absl::StrCat(
+        "solved dynamic expressions do not map to exactly one shared dynamic "
+        "variable id: ",
+        absl::StrJoin(expr_summaries, " | "));
+  } else {
+    result.can_run = false;
+    result.diagnostic = absl::StrCat(
+        "Failed to recover a unique XLA dynamic batch size from runtime "
+        "input expressions. solved_expressions=[",
+        absl::StrJoin(expr_summaries, " | "), "] all_values={",
+        absl::StrJoin(dyn_vals, ", "), "}");
+  }
+  return result;
+}
+
+
+}  // namespace xla_ops_internal
+
+namespace {
+
+using xla_ops_internal::AnalyzeIgnoredDynamicArgumentOccurrences;
+using xla_ops_internal::BuildStaticCompilationArguments;
+using xla_ops_internal::DynamicBatchResolutionResult;
+using xla_ops_internal::DynamicSolveFilterDecision;
+using xla_ops_internal::GetRuntimeInputIndex;
+using xla_ops_internal::StripIgnoredDynamicArgumentOccurrences;
+
+DynamicBatchResolutionResult ResolveDynamicBatchSizeFromRuntimeInputs(
+    OpKernelContext* ctx, const XlaCompiler::CompilationResult& comp_result,
+    int num_constant_args, bool log_solves, absl::string_view cluster_name,
+    const NodeDef& op_def) {
+  DynamicBatchResolutionResult result;
+  if (!GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes) {
+    return result;
+  }
+
+  auto get_runtime_input_index = [&](int i) {
+    int input_idx =
+        GetRuntimeInputIndex(comp_result.input_mapping, i, num_constant_args,
+                             /*constants_omitted=*/op_def.op() == "_XlaRun");
+    const bool is_runtime_key_input =
+        input_idx >= 0 && input_idx < ctx->num_inputs() &&
+        ctx->input_dtype(input_idx) == DT_STRING &&
+        input_idx == op_def.input_size() - 1;
+    return is_runtime_key_input ? -1 : input_idx;
+  };
+  auto get_runtime_input_shape = [&](int i) -> const TensorShape& {
+    return ctx->input(i).shape();
+  };
+  auto get_runtime_input_name = [&](int i) -> absl::string_view {
+    return i < op_def.input_size() ? op_def.input(i) : "<unknown>";
+  };
+
+  result = xla_ops_internal::ResolveDynamicBatchSizeFromRuntimeShapes(
+      comp_result.xla_input_shapes, ctx->num_inputs(),
+      get_runtime_input_index, get_runtime_input_shape,
+      get_runtime_input_name, log_solves, cluster_name);
+  if (result.has_batch_size || !result.can_run) {
+    return result;
+  }
+
+  BatchSizeResource* bsr = nullptr;
+  ScopedStepContainer* step_container = ctx->step_container();
+  absl::Status st = step_container->Lookup<BatchSizeResource>(
+      ctx->resource_manager(), BatchSizeResourceName, &bsr);
+
+  if (st.ok()) {
+    CHECK(bsr != nullptr);
+    result.has_batch_size = true;
+    result.batch_size = bsr->GetBatchSize();
+    result.diagnostic = absl::StrCat(
+        "BatchSizeResource fallback value=", result.batch_size);
+    bsr->Unref();
+  } else if (IsNotFound(st)) {
+    result.batch_size_resource_not_found = true;
+  } else {
+    result.can_run = false;
+    result.diagnostic = st.ToString();
+  }
+  return result;
+}
+
 se::Stream* GetStream(OpKernelContext* ctx) {
   return ctx->op_device_context() ? ctx->op_device_context()->stream()
                                   : nullptr;
@@ -384,6 +959,7 @@ absl::Status CompileToLocalExecutable(
     const XlaPlatformInfo& platform_info,
     const std::vector<XlaCompiler::Argument>& args,
     DeviceCompileMode compile_mode, bool may_alias_resource_update,
+    bool force_static_shapes, bool* dynamic_solve_conflict,
     xla::LocalClient** client,
     const XlaCompiler::CompilationResult** compilation_result,
     xla::LocalExecutable** executable) {
@@ -427,9 +1003,415 @@ absl::Status CompileToLocalExecutable(
   XlaCompiler::CompileOptions compile_options =
       GenerateCompileOptions(has_ref_vars, may_alias_resource_update);
 
-  return xla_device_compiler->CompileIfNeeded(
-      options, function, args, compile_options, compile_mode, profiler,
-      compilation_result, executable);
+  if (dynamic_solve_conflict != nullptr) {
+    *dynamic_solve_conflict = false;
+  }
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
+  if (flags->tf_xla_enable_dynamic_sizes && !force_static_shapes) {
+    // Rewriting the argument with expressions if they have dynamic
+    // dimension, detecting dynamic dimension via either _dynamic_dim or the
+    // inferred-output-shapes attr attached during encapsulation.
+    std::vector<XlaCompiler::Argument> norm_args(args.begin(), args.end());
+    int64_t filled_batch = 0;
+    bool saw_dynamic_dim_value = false;
+    // Only supporting one dynamic dimension.
+    bool has_multiple_dynamic_dim_values = false;
+    int64_t dynamic_dim_value = 0;
+    XlaBatchMatcher* xla_batch_matcher =
+        xla_device_compiler->xla_batch_matcher();
+    std::optional<xla::DExpr> dynamic_dim_expr;
+    std::optional<xla::DExpr> shared_dynamic_subexpr;
+    auto normalize_dynamic_expr =
+        [&](xla::DExpr expr, absl::string_view context = "") {
+          if (!expr || !expr->is_dynamic()) {
+            return expr;
+          }
+          if (!shared_dynamic_subexpr.has_value()) {
+            shared_dynamic_subexpr =
+                expr.find_smallest_subexpression_covering_all_variables();
+            VLOG(1) << "Using shared dynamic subexpression "
+                      << DExprToString(*shared_dynamic_subexpr)
+                      << " for XLA dynamic input normalization. context="
+                      << context;
+          }
+          xla::DExpr normalized_expr =
+              expr.replace_subexpression(*shared_dynamic_subexpr,
+                                         xla::DExpr::Var(1))
+                  .simplify();
+          VLOG(1) << "Rewriting dynamic input expression "
+                    << DExprToString(expr)
+                    << " to shared-core form "
+                    << DExprToString(normalized_expr)
+                    << " before XLA compilation. context=" << context;
+          return normalized_expr;
+        };
+    auto maybe_attach_shape_contents_from_attrs =
+        [&](int arg_index, const auto& attr_map,
+            const std::string& node_name) {
+          auto& arg = norm_args[arg_index];
+          if (arg.kind != XlaCompiler::Argument::kConstant) {
+            return;
+          }
+
+          auto inferred_shape_it = attr_map.find("user_inferred_shape");
+          auto inferred_contents_it =
+              attr_map.find(kUserInferredValueContentsAttrName);
+          bool has_dynamic = false;
+          auto has_dynamic_it = attr_map.find("has_dynamic");
+          if (has_dynamic_it != attr_map.end()) {
+            has_dynamic = has_dynamic_it->second.b();
+          }
+
+          if (inferred_contents_it == attr_map.end() &&
+              (!has_dynamic || inferred_shape_it == attr_map.end())) {
+            return;
+          }
+
+          TensorShapeProto inferred_shape_proto;
+          if (inferred_contents_it != attr_map.end()) {
+            if (!inferred_shape_proto.ParseFromString(
+                    inferred_contents_it->second.s())) {
+              return;
+            }
+          } else {
+            inferred_shape_proto = inferred_shape_it->second.shape();
+          }
+
+          TensorShape inferred_shape(inferred_shape_proto);
+          if (!((TensorShapeUtils::IsVector(arg.constant_value.shape()) &&
+                 arg.constant_value.NumElements() == inferred_shape.dims()) ||
+                (TensorShapeUtils::IsScalar(arg.constant_value.shape()) &&
+                 inferred_shape.dims() == 1))) {
+            return;
+          }
+
+          arg.constant_value_expressions.clear();
+          arg.constant_value_expressions.reserve(inferred_shape.dims());
+          for (int64_t i = 0; i < inferred_shape.dims(); ++i) {
+            xla::ExpressionProto expr;
+            const xla::DExpr& dim_expr = inferred_shape.get_expression(i);
+            if (dim_expr && dim_expr->is_dynamic()) {
+              xla::DExpr normalized_expr = normalize_dynamic_expr(
+                  dim_expr,
+                  absl::StrCat("const_arg=", arg_index, " node=", node_name,
+                               " dim=", i));
+              normalized_expr->to_proto(&expr);
+            } else if (arg.constant_value.dtype() == DT_INT32) {
+              expr.set_constant_value(arg.constant_value.flat<int32>()(i));
+            } else if (arg.constant_value.dtype() == DT_INT64) {
+              expr.set_constant_value(arg.constant_value.flat<int64_t>()(i));
+            } else {
+              arg.constant_value_expressions.clear();
+              return;
+            }
+            arg.constant_value_expressions.push_back(std::move(expr));
+          }
+        };
+    auto record_dynamic_dim_value = [&](int64_t dim_size, xla::DExpr expr) {
+      if (!saw_dynamic_dim_value) {
+        saw_dynamic_dim_value = true;
+        dynamic_dim_value = dim_size;
+        dynamic_dim_expr = std::move(expr);
+        return;
+      }
+      if (dynamic_dim_value != dim_size) {
+        has_multiple_dynamic_dim_values = true;
+      }
+    };
+    if (options.flib_def != nullptr) {
+      const FunctionDef* fdef = options.flib_def->Find(function.name());
+      if (fdef != nullptr) {
+        for (const auto& kv : fdef->arg_attr()) {
+          int arg_index = kv.first;
+          const auto& attr_map = kv.second.attr();
+          const std::string& node_name =
+              fdef->signature().input_arg(arg_index).name();
+
+          auto shape_derived_attr = attr_map.find(kXlaShapeDerivedAttrName);
+          if (shape_derived_attr != attr_map.end()) {
+            VLOG(1) << "XlaCompileOp retrieved shape-derived marker for arg "
+                    << arg_index << " node=" << node_name;
+          }
+          maybe_attach_shape_contents_from_attrs(arg_index, attr_map, node_name);
+
+          // Special case for _dynamic_dim...
+          auto dyn_dim_attr = attr_map.find("_dynamic_dim");
+          if (dyn_dim_attr != attr_map.end()) {
+            TensorShape& shp =
+                std::get<TensorShape>(norm_args[arg_index].shape);
+            const AttrValue& v = dyn_dim_attr->second;
+            int64_t idx = v.i();
+            record_dynamic_dim_value(shp.dim_size(idx), xla::DExpr::Var(1));
+            if (!filled_batch && xla_batch_matcher) {
+              filled_batch =
+                  xla_batch_matcher->get_xla_compile_batch(function.name(), shp.dim_size(idx));
+            }
+
+            std::vector<xla::DExpr> dyn_exprs;
+            for (int d : shp.dim_sizes()) {
+              dyn_exprs.push_back(xla::DExpr::Const(d));
+            }
+            dyn_exprs[idx] = *dynamic_dim_expr;
+            shp.set_expressions(std::move(dyn_exprs));
+            continue;
+          }
+          auto it = attr_map.find(kXlaInferredOutputShapesAttrName);
+          if (it == attr_map.end()) continue;
+
+          const TensorShapeProto& proto = it->second.list().shape(0);
+          const auto& exp = proto.expressions();
+          TensorShape& shp = std::get<TensorShape>(norm_args[arg_index].shape);
+          if (!filled_batch && xla_batch_matcher) {
+            for (int idx = 0; idx < exp.size(); ++idx) {
+              // Look for dynamic expression. If found then compute padding
+              // value and exit loop.
+              auto e = normalize_dynamic_expr(DimExprFromProto(exp[idx]));
+              if (e->is_dynamic()) {
+                VLOG(1) << "Calling dynamic expression solve for compile "
+                          << "argument " << arg_index << " dimension " << idx
+                          << " expr=" << DExprToString(e)
+                          << " target_size=" << shp.dim_size(idx);
+                std::optional<int64_t> solved_value =
+                    e->solve(shp.dim_size(idx));
+                VLOG(1) << "Dynamic expression solve for compile argument "
+                          << arg_index << " dimension " << idx << " returned "
+                          << (solved_value.has_value()
+                                  ? std::to_string(*solved_value)
+                                  : std::string("<none>"));
+                int64_t var_value;
+                if (!solved_value.has_value()) {
+                  LOG(WARNING)
+                      << "Failed to solve dynamic dimension for argument "
+                      << arg_index << " dim " << idx << " with size "
+                      << shp.dim_size(idx)
+                      << "; falling back to original dimension size.";
+                  var_value = shp.dim_size(idx);
+                } else {
+                  var_value = *solved_value;
+                  VLOG(1) << "Solved dynamic dimension from "
+                          << shp.dim_size(idx) << " to " << var_value;
+                }
+                record_dynamic_dim_value(var_value, e);
+                VLOG(1) << "Filled batch for function " << function.name()
+                          << " with value " << var_value;
+                filled_batch =
+                    xla_batch_matcher->get_xla_compile_batch(function.name(), var_value);
+                break;
+              }
+            }
+          }
+
+          std::vector<xla::DExpr> dyn_exprs;
+          for (int d : shp.dim_sizes()) {
+            dyn_exprs.push_back(xla::DExpr::Const(d));
+          }
+          for (int j = 0; j < exp.size(); ++j) {
+            auto e = DimExprFromProto(exp[j]);
+            if (e->is_dynamic()) {
+              e = normalize_dynamic_expr(
+                  e, absl::StrCat("arg=", arg_index, " dim=", j,
+                                  " input_shape"));
+              dyn_exprs[j] = e;
+            }
+          }
+          shp.set_expressions(std::move(dyn_exprs));
+        }
+      }
+    }
+
+    DynamicSolveFilterDecision solve_filter_decision =
+        AnalyzeIgnoredDynamicArgumentOccurrences(norm_args);
+    if (!solve_filter_decision.can_run) {
+      if (dynamic_solve_conflict != nullptr) {
+        *dynamic_solve_conflict = true;
+      }
+      if (compile_mode == DeviceCompileMode::kLazy) {
+        return errors::Unimplemented(solve_filter_decision.diagnostic);
+      }
+      return errors::InvalidArgument(solve_filter_decision.diagnostic);
+    }
+    if (!solve_filter_decision.ignored_occurrences.empty()) {
+      LOG(WARNING) << solve_filter_decision.diagnostic;
+      StripIgnoredDynamicArgumentOccurrences(
+          solve_filter_decision.ignored_occurrences, &norm_args);
+      saw_dynamic_dim_value = false;
+      has_multiple_dynamic_dim_values = false;
+      dynamic_dim_value = 0;
+      dynamic_dim_expr.reset();
+      filled_batch = 0;
+      for (int arg_index = 0; arg_index < norm_args.size(); ++arg_index) {
+        if (!absl::holds_alternative<TensorShape>(norm_args[arg_index].shape)) {
+          continue;
+        }
+        TensorShape& shape = std::get<TensorShape>(norm_args[arg_index].shape);
+        for (int dim = 0; dim < shape.get_expressions().size(); ++dim) {
+          xla::DExpr expr = shape.get_expression(dim);
+          if (!(expr && expr->is_dynamic())) {
+            continue;
+          }
+          xla::DExpr simplified_expr = expr.simplify();
+          std::optional<int64_t> solved_value =
+              simplified_expr->solve(shape.dim_size(dim));
+          if (!solved_value.has_value()) {
+            continue;
+          }
+          record_dynamic_dim_value(*solved_value, simplified_expr);
+          if (!filled_batch && xla_batch_matcher) {
+            filled_batch =
+                xla_batch_matcher->get_xla_compile_batch(*solved_value);
+          }
+        }
+      }
+    }
+
+    struct SaveOldVar {
+      int arg_index;
+      int64_t dyn_dim;
+      int64_t old_value;
+    };
+    std::vector<SaveOldVar> old_vars;
+    auto maybe_rewrite_scalar_constant = [&](int arg_index) {
+      if (!saw_dynamic_dim_value || has_multiple_dynamic_dim_values) {
+        return;
+      }
+
+      auto& arg = norm_args[arg_index];
+      if (arg.kind != XlaCompiler::Argument::kConstant) {
+        return;
+      }
+
+      const bool is_scalar = TensorShapeUtils::IsScalar(arg.constant_value.shape());
+      const bool is_vector = TensorShapeUtils::IsVector(arg.constant_value.shape()) &&
+                             arg.constant_value.NumElements() > 0;
+      if (!is_scalar && !is_vector) {
+        return;
+      }
+
+      auto set_constant_contents = [&]<typename T>(int rewrite_index) {
+        arg.constant_value_expressions.clear();
+        const int64_t num_elements = arg.constant_value.NumElements();
+        arg.constant_value_expressions.reserve(num_elements);
+        for (int64_t i = 0; i < num_elements; ++i) {
+          xla::ExpressionProto expr;
+          if (i == rewrite_index) {
+            dynamic_dim_expr->to_proto(&expr);
+          } else {
+            expr.set_constant_value(arg.constant_value.flat<T>()(i));
+          }
+          arg.constant_value_expressions.push_back(std::move(expr));
+        }
+      };
+
+      if (arg.constant_value.dtype() == DT_INT32) {
+        auto flat = arg.constant_value.flat<int32>();
+        int rewrite_index = -1;
+        // Heuristic: rewrite only scalar constants or shape-like int vectors.
+        // In practice we expect at most one entry to match the observed
+        // runtime batch size, so rewrite the first matching entry.
+        for (int i = 0; i < arg.constant_value.NumElements(); ++i) {
+          if (flat(i) == dynamic_dim_value) {
+            rewrite_index = i;
+            break;
+          }
+        }
+        if (rewrite_index >= 0) {
+          arg.constant_value = tensor::DeepCopy(arg.constant_value);
+          auto mutable_flat = arg.constant_value.flat<int32>();
+          VLOG(1) << "XlaCompileOp int32 constant arg " << arg_index
+                  << " index " << rewrite_index
+                  << " matches dynamic_dim_value=" << dynamic_dim_value;
+          mutable_flat(rewrite_index) = filled_batch;
+          set_constant_contents.template operator()<int32>(rewrite_index);
+        }
+      } else if (arg.constant_value.dtype() == DT_INT64) {
+        auto flat = arg.constant_value.flat<int64_t>();
+        int rewrite_index = -1;
+        // Same heuristic for int64 scalar constants or shape-like vectors.
+        for (int i = 0; i < arg.constant_value.NumElements(); ++i) {
+          if (flat(i) == dynamic_dim_value) {
+            rewrite_index = i;
+            break;
+          }
+        }
+        if (rewrite_index >= 0) {
+          arg.constant_value = tensor::DeepCopy(arg.constant_value);
+          auto mutable_flat = arg.constant_value.flat<int64_t>();
+          VLOG(1) << "XlaCompileOp int64 constant arg " << arg_index
+                  << " index " << rewrite_index
+                  << " matches dynamic_dim_value=" << dynamic_dim_value;
+          mutable_flat(rewrite_index) = filled_batch;
+          set_constant_contents.template operator()<int64_t>(rewrite_index);
+        }
+      }
+    };
+    // We rewrite only dynamic dimensions to the padded compile batch and then
+    // restore the original runtime sizes after compilation. Some scalar
+    // constants are actually runtime batch sizes folded by earlier TF passes,
+    // so rewrite only those that match the detected dynamic runtime value.
+    // Scalar constants are deep-copied before rewrite so the change stays
+    // local to norm_args and does not require restoration.
+    if (filled_batch) {
+      for (int i = 0; i < norm_args.size(); ++i) {
+        TensorShape& shp = std::get<TensorShape>(norm_args[i].shape);
+        for (int j = 0; j < shp.get_expressions().size(); ++j) {
+          auto e = shp.get_expression(j);
+          if (e && e->is_dynamic()) {
+            int64_t old = shp.dim_size(j);
+            old_vars.push_back({i, j, old});
+            xla::DExpr padded_expr = xla::DExpr::Const(filled_batch);
+            const std::set<int> ids = e->get_all_ids();
+            if (ids.size() != 1) {
+              return errors::InvalidArgument(
+                  "Dynamic shape padding expected exactly one dynamic "
+                  "variable for argument ",
+                  i, ", dimension ", j, ", but found ", ids.size(),
+                  " variables in expression ", DExprToString(e));
+            }
+            const int substitute_var_id = *ids.begin();
+            VLOG(1) << "Calling dynamic expression substitute for compile "
+                      << "argument " << i << " dimension " << j
+                      << " expr=" << DExprToString(e)
+                      << " substitute Var(" << substitute_var_id
+                      << ")=" << filled_batch;
+            xla::DExpr subst_expr =
+                e.substitute(substitute_var_id, padded_expr).simplify();
+            VLOG(1) << "Dynamic expression substitute for compile argument "
+                      << i << " dimension " << j
+                      << " returned " << DExprToString(subst_expr);
+            if (!subst_expr->is_constant()) {
+              return errors::InvalidArgument(
+                  "Dynamic shape padding substitution did not produce an "
+                  "integer constant for argument ",
+                  i, ", dimension ", j, ": ", DExprToString(subst_expr));
+            }
+            int64_t new_dim = subst_expr->get_val();
+            if (new_dim >= 0) {
+              shp.set_dim(j, new_dim);
+              // Necessary because set_dim removes the expression:
+              shp.set_expression(j, e);
+            }
+          }
+        }
+        maybe_rewrite_scalar_constant(i);
+      }
+    }
+    auto status = xla_device_compiler->CompileIfNeeded(
+        options, function, norm_args, compile_options, compile_mode, profiler,
+        compilation_result, executable);
+    // Restore the original runtime dimensions after compilation.
+    if (filled_batch) {
+      for (const auto& old_var : old_vars) {
+        TensorShape& shp =
+            std::get<TensorShape>(norm_args[old_var.arg_index].shape);
+        shp.set_dim(old_var.dyn_dim, old_var.old_value);
+      }
+    }
+    return status;
+  } else {
+    return xla_device_compiler->CompileIfNeeded(
+        options, function, args, compile_options, compile_mode, profiler,
+        compilation_result, executable);
+  }
 }
 
 absl::Status GetUpdatedVariables(
@@ -582,11 +1564,25 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
     return;
   }
 
+  bool dynamic_solve_conflict = false;
   absl::Status status = CompileToLocalExecutable(
       ctx, function_, /*has_ref_vars=*/has_ref_vars_, platform_info_,
       xla_compiler_args, DeviceCompileMode::kStrict,
-      /*may_alias_resource_update=*/true, &client, &compilation_result,
+      /*may_alias_resource_update=*/true, /*force_static_shapes=*/false,
+      &dynamic_solve_conflict, &client, &compilation_result,
       &executable);
+  if (dynamic_solve_conflict) {
+    LOG(WARNING) << "Retrying XLA cluster " << function_.name()
+                 << " with concrete static argument shapes";
+    std::vector<XlaCompiler::Argument> static_args =
+        BuildStaticCompilationArguments(xla_compiler_args);
+    status = CompileToLocalExecutable(
+        ctx, function_, /*has_ref_vars=*/has_ref_vars_, platform_info_,
+        static_args, DeviceCompileMode::kStrict,
+        /*may_alias_resource_update=*/true, /*force_static_shapes=*/true,
+        /*dynamic_solve_conflict=*/nullptr, &client, &compilation_result,
+        &executable);
+  }
   OP_REQUIRES_OK_ASYNC(ctx, status, done);
 
   // Continuation of the execution, may be run in a different thread.
@@ -775,6 +1771,27 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) {
       GetXlaOpsCommonFlags()
           ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice(
               platform_info_.device_type());
+  std::vector<XlaCompiler::Argument> compiler_args;
+  auto compile_static_fallback = [&]() -> absl::Status {
+    std::vector<XlaCompiler::Argument> static_args =
+        BuildStaticCompilationArguments(compiler_args);
+    kernel = nullptr;
+    executable = nullptr;
+    pjrt_executable = nullptr;
+    LOG(WARNING) << "Retrying XLA cluster " << function_.name()
+                 << " with concrete static argument shapes";
+    if (use_pjrt) {
+      return CompileToPjRtLoadedExecutable(
+          *ctx, platform_info_, function_, static_args, compile_mode,
+          has_ref_vars_, /*may_alias_resource_update=*/false, &kernel,
+          &pjrt_client, &pjrt_executable);
+    }
+    return CompileToLocalExecutable(
+        ctx, function_, has_ref_vars_, platform_info_, static_args,
+        compile_mode, /*may_alias_resource_update=*/false,
+        /*force_static_shapes=*/true,
+        /*dynamic_solve_conflict=*/nullptr, &client, &kernel, &executable);
+  };
 
   if (GetXlaOpsCommonFlags()->tf_xla_always_defer_compilation ||
       cannot_compile_cluster) {
@@ -783,13 +1800,14 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) {
     auto args_and_variables_snapshot = GetXlaCompilerArgsAndSnapshotVariables(
         resources_, constants_, inputs, ctx);
     OP_REQUIRES_OK(ctx, args_and_variables_snapshot.status());
-    const std::vector<XlaCompiler::Argument>& args =
-        args_and_variables_snapshot->first;
+    compiler_args = std::move(args_and_variables_snapshot->first);
+    const std::vector<XlaCompiler::Argument>& args = compiler_args;
     variables_snapshot = std::move(args_and_variables_snapshot->second);
 
     // Do not alias resource updates as locking variables in XlaCompile and
     // unlocking them in XlaRun may lead to deadlocks.
     absl::Status status;
+    bool dynamic_solve_conflict = false;
     if (use_pjrt) {
       VLOG(2) << "Using PJRT for compilation. Function name: "
               << function_.name();
@@ -800,16 +1818,30 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) {
     } else {
       status = CompileToLocalExecutable(
           ctx, function_, has_ref_vars_, platform_info_, args, compile_mode,
-          /*may_alias_resource_update=*/false, &client, &kernel, &executable);
+          /*may_alias_resource_update=*/false, /*force_static_shapes=*/false,
+          &dynamic_solve_conflict, &client, &kernel, &executable);
+    }
+    if (dynamic_solve_conflict) {
+      status = compile_static_fallback();
     }
+
     if (compile_mode != DeviceCompileMode::kLazy ||
         status.code() != error::UNIMPLEMENTED) {
-      OP_REQUIRES_OK(ctx, status);
+      if ((status != OkStatus()) &&
+          (status.code() != error::UNIMPLEMENTED) &&
+        (compile_mode == DeviceCompileMode::kLazy)) {
+        // We set the error to error::UNIMPLEMENTED so it falls in the
+        // conditions of the if to fall back to TensorFlow function call
+        status = tensorflow::errors::Unimplemented(status.ToString());
+      } else {
+        OP_REQUIRES_OK(ctx, status);
+      }
     }
 
     if (status.code() == error::UNIMPLEMENTED) {
-      LOG(WARNING) << "Compilation failed:" << status
-                   << ".  Falling back to TF function call.";
+      LOG(WARNING) << "[HUAWEI] Compilation of the cluster failed with:";
+      LOG(WARNING) << "[HUAWEI] " << status;
+      LOG(WARNING) << "[HUAWEI] Falling back to TF function call.\n";
 
       BroadcastOptimizationRemark(
           XlaOptimizationRemark::UNIMPLEMENTED_OPERATION, status.ToString())
@@ -817,10 +1849,41 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) {
       executable = nullptr;
       pjrt_executable = nullptr;
       mutex_lock guard(cannot_compile_cluster_mu_);
+      // TODO: decide if we want to set this flag to true, as we may want to
+      // allow the cluster to try to compile again later in time.
       cannot_compile_cluster_ = true;
     }
   }
 
+  if ((executable || pjrt_executable) && kernel != nullptr &&
+      GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes) {
+    DynamicBatchResolutionResult resolution =
+        ResolveDynamicBatchSizeFromRuntimeInputs(ctx, *kernel, constants_.size(),
+                                                 /*log_solves=*/true,
+                                                 function_.name(), def());
+    if (!resolution.can_run) {
+      const std::string error_message = absl::StrCat(
+          "Rejecting XLA cluster at compile time because dynamic expressions "
+          "cannot be solved consistently for this request. ",
+          resolution.diagnostic);
+      LOG(WARNING) << error_message;
+      absl::Status static_status = compile_static_fallback();
+      if (!static_status.ok()) {
+        if (must_compile_) {
+          OP_REQUIRES_OK(ctx, static_status);
+        }
+        LOG(WARNING) << "Static XLA fallback failed for cluster "
+                     << function_.name() << ": " << static_status;
+        executable = nullptr;
+        pjrt_executable = nullptr;
+        kernel = nullptr;
+      } else {
+        LOG(INFO) << "Using static XLA fallback for cluster "
+                  << function_.name();
+      }
+    }
+  }
+
   AllocatorAttributes host_alloc_attrs;
   host_alloc_attrs.set_gpu_compatible(true);
   host_alloc_attrs.set_on_host(true);
@@ -871,7 +1934,6 @@ XlaRunOp::XlaRunOp(OpKernelConstruction* ctx)
 void XlaRunOp::Compute(OpKernelContext* ctx) {
   VLOG(3) << "XlaRunOp " << def().name();
   Tensor key_tensor = ctx->input(ctx->num_inputs() - 1);
-
   bool use_pjrt =
       GetXlaOpsCommonFlags()
           ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice(
@@ -881,6 +1943,14 @@ void XlaRunOp::Compute(OpKernelContext* ctx) {
     const PjRtExecutableClosureStore::KeyT& key = key_tensor.flat<tstring>()(0);
     PjRtExecutableClosure closure =
         PjRtExecutableClosureStore::Global()->Consume(key);
+    const std::string cluster_name =
+        closure.compilation_result() != nullptr &&
+                closure.compilation_result()->computation != nullptr
+            ? closure.compilation_result()->computation->name()
+            : std::string("<unknown>");
+    VLOG(1) << "Entering XLA cluster: cluster=" << cluster_name
+              << " op=" << def().name() << " key=" << key
+              << " mode=pjrt step_id=" << ctx->step_id();
 
     // Fetch inputs from the OpKernelContext. Inputs are the same as the ones
     // for XlaCompile, except that the must-be-constant inputs that appear in
@@ -917,6 +1987,14 @@ void XlaRunOp::Compute(OpKernelContext* ctx) {
 
   XlaExecutableClosure closure =
       XlaExecutableClosureStore::Global()->Consume(key);
+  const std::string cluster_name =
+      closure.compilation_result() != nullptr &&
+              closure.compilation_result()->computation != nullptr
+          ? closure.compilation_result()->computation->name()
+          : std::string("<unknown>");
+  VLOG(1) << "Entering XLA cluster: cluster=" << cluster_name
+            << " op=" << def().name() << " key=" << key
+            << " mode=local step_id=" << ctx->step_id();
   std::shared_ptr<se::DeviceMemoryAllocator> allocator =
       GetAllocator(ctx->device(), GetStream(ctx), platform_info_);
   XlaComputationLaunchContext launch_context =
@@ -953,6 +2031,36 @@ void XlaRunOp::Compute(OpKernelContext* ctx) {
 
   xla::ExecutableRunOptions run_options;
 
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    DynamicBatchResolutionResult batch_resolution =
+        ResolveDynamicBatchSizeFromRuntimeInputs(
+            ctx, *closure.compilation_result(), closure.num_constant_args(),
+            /*log_solves=*/true, cluster_name, def());
+    bool is_set = false;
+    if (!batch_resolution.can_run) {
+      LOG(ERROR) << batch_resolution.diagnostic;
+      ctx->CtxFailure(errors::InvalidArgument(batch_resolution.diagnostic));
+      return;
+    }
+    if (batch_resolution.has_batch_size) {
+      VLOG(1) << "Setting run_options.batch_size "
+                << batch_resolution.diagnostic;
+      run_options.set_batch_size(batch_resolution.batch_size);
+      is_set = true;
+    }
+    if (!is_set) {
+      LOG(WARNING) << "Entering XLA cluster without run_options.batch_size "
+                   << "being set because " << batch_resolution.diagnostic
+                   << ". op=" << def().name() << " closure_key=" << key
+                   << " step_id=" << ctx->step_id()
+                   << " current_run_options_batch_size="
+                   << run_options.batch_size()
+                   << " batch_size_resource_not_found="
+                   << batch_resolution.batch_size_resource_not_found;
+    }
+  }
+
   // Host callbacks used for HLO send/recv.
   xla::SendDeviceMemoryFunction send_function =
       GetSendDeviceMemoryFunction(ctx, key);
@@ -981,7 +2089,8 @@ void XlaRunOp::Compute(OpKernelContext* ctx) {
       launch_context.PopulateOutputs(
           ctx, closure.compilation_result(), execution_output->ConsumeResult(),
           /*missing_ctx_input_prefix=*/closure.num_constant_args(),
-          absl::MakeSpan(*variable_infos), input_output_alias, snapshot_ptrs));
+          absl::MakeSpan(*variable_infos), input_output_alias, snapshot_ptrs,
+          &run_options));
 }
 
 XlaMergeOp::XlaMergeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
diff --git a/tensorflow/compiler/jit/kernels/xla_ops_internal.h b/tensorflow/compiler/jit/kernels/xla_ops_internal.h
new file mode 100644
index 00000000000000..546bc46d93534a
--- /dev/null
+++ b/tensorflow/compiler/jit/kernels/xla_ops_internal.h
@@ -0,0 +1,84 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_INTERNAL_H_
+#define TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_INTERNAL_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/functional/function_ref.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/xla_argument.h"
+
+namespace tensorflow {
+namespace xla_ops_internal {
+
+struct IgnoredDynamicArgumentOccurrence {
+  enum class Source {
+    kShapeDimension,
+    kConstantValueElement,
+  };
+
+  Source source;
+  int arg_index;
+  int dim_or_index;
+  int64_t observed_value;
+  int64_t solved_value;
+  std::string expr;
+};
+
+struct DynamicSolveFilterDecision {
+  bool can_run = true;
+  std::string diagnostic;
+  std::vector<IgnoredDynamicArgumentOccurrence> ignored_occurrences;
+};
+
+struct DynamicBatchResolutionResult {
+  bool can_run = true;
+  bool has_batch_size = false;
+  int64_t batch_size = 0;
+  bool batch_size_resource_not_found = false;
+  std::string diagnostic;
+};
+
+int GetRuntimeInputIndex(absl::Span<const int> input_mapping,
+                         int xla_input_index, int num_constant_args,
+                         bool constants_omitted);
+
+DynamicBatchResolutionResult ResolveDynamicBatchSizeFromRuntimeShapes(
+    absl::Span<const xla::Shape> xla_input_shapes,
+    int runtime_input_count,
+    absl::FunctionRef<int(int)> get_runtime_input_index,
+    absl::FunctionRef<const TensorShape&(int)> get_runtime_input_shape,
+    absl::FunctionRef<absl::string_view(int)> get_runtime_input_name,
+    bool log_solves = false, absl::string_view cluster_name = {});
+
+DynamicSolveFilterDecision AnalyzeIgnoredDynamicArgumentOccurrences(
+    absl::Span<const XlaArgument> args);
+
+std::vector<XlaArgument> BuildStaticCompilationArguments(
+    absl::Span<const XlaArgument> args);
+
+void StripIgnoredDynamicArgumentOccurrences(
+    const std::vector<IgnoredDynamicArgumentOccurrence>& ignored_occurrences,
+    std::vector<XlaArgument>* args);
+
+}  // namespace xla_ops_internal
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_INTERNAL_H_
diff --git a/tensorflow/compiler/jit/kernels/xla_ops_test.cc b/tensorflow/compiler/jit/kernels/xla_ops_test.cc
new file mode 100644
index 00000000000000..bcea59771626df
--- /dev/null
+++ b/tensorflow/compiler/jit/kernels/xla_ops_test.cc
@@ -0,0 +1,282 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/kernels/xla_ops_internal.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "xla/shape_expr.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace xla_ops_internal {
+namespace {
+
+TEST(XlaInputMappingTest, UsesCompilationInputMapping) {
+  const std::vector<int> input_mapping = {1, 3};
+
+  EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 0, /*num_constant_args=*/1,
+                                 /*constants_omitted=*/false),
+            1);
+  EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 1, /*num_constant_args=*/1,
+                                 /*constants_omitted=*/false),
+            3);
+  EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 0, /*num_constant_args=*/1,
+                                 /*constants_omitted=*/true),
+            0);
+  EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 1, /*num_constant_args=*/1,
+                                 /*constants_omitted=*/true),
+            2);
+  EXPECT_EQ(GetRuntimeInputIndex(input_mapping, 2, /*num_constant_args=*/1,
+                                 /*constants_omitted=*/true),
+            -1);
+}
+
+XlaArgument MakeDynamicArgument(int64_t observed_size,
+                                xla::DExpr expr = xla::DExpr::Var(1)) {
+  XlaArgument arg;
+  arg.kind = XlaArgument::kParameter;
+  arg.type = DT_FLOAT;
+  TensorShape shape({observed_size, 16});
+  shape.set_expression(0, std::move(expr));
+  shape.set_expression(1, xla::DExpr::Const(16));
+  arg.shape = shape;
+  return arg;
+}
+
+xla::Shape MakeDynamicXlaInput(xla::DExpr leading_expr) {
+  xla::Shape shape;
+  shape.set_element_type(xla::F32);
+  shape.add_dimensions(1024, /*is_dynamic=*/true,
+                       std::move(leading_expr));
+  shape.add_dimensions(16, /*is_dynamic=*/false, xla::DExpr::Const(16));
+  return shape;
+}
+
+DynamicBatchResolutionResult ResolveForTest(
+    absl::Span<const xla::Shape> xla_shapes,
+    absl::Span<const TensorShape> runtime_shapes,
+    absl::Span<const int> runtime_input_indices,
+    absl::Span<const std::string> runtime_input_names) {
+  return ResolveDynamicBatchSizeFromRuntimeShapes(
+      xla_shapes, runtime_shapes.size(),
+      [&](int i) { return runtime_input_indices[i]; },
+      [&](int i) -> const TensorShape& { return runtime_shapes[i]; },
+      [&](int i) -> absl::string_view { return runtime_input_names[i]; });
+}
+
+TEST(DynamicRuntimeSolveTest, UsesMappedRuntimeShapesToSolveExpressions) {
+  const xla::DExpr variable = xla::DExpr::Var(1);
+  const std::vector<xla::Shape> xla_shapes = {
+      MakeDynamicXlaInput(variable),
+      MakeDynamicXlaInput(variable / 2),
+  };
+  // XLA parameters are deliberately ordered differently from runtime inputs.
+  const std::vector<TensorShape> runtime_shapes = {
+      TensorShape({120, 16}),
+      TensorShape({240, 16}),
+  };
+  const std::vector<int> runtime_input_indices = {1, 0};
+  const std::vector<std::string> runtime_input_names = {"half", "full"};
+
+  DynamicBatchResolutionResult result = ResolveForTest(
+      xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names);
+
+  EXPECT_TRUE(result.can_run);
+  ASSERT_TRUE(result.has_batch_size);
+  EXPECT_EQ(result.batch_size, 240);
+}
+
+TEST(DynamicRuntimeSolveTest,
+     PrefersNonSingletonEvidenceAcrossEquivalentExpressions) {
+  const xla::DExpr variable = xla::DExpr::Var(1);
+  const std::vector<xla::Shape> xla_shapes = {
+      MakeDynamicXlaInput(variable),
+      MakeDynamicXlaInput(variable + 1),
+  };
+  // The first occurrence may be a broadcast singleton. A + 1 = 241 gives
+  // the stronger runtime evidence A = 240.
+  const std::vector<TensorShape> runtime_shapes = {
+      TensorShape({1, 16}),
+      TensorShape({241, 16}),
+  };
+  const std::vector<int> runtime_input_indices = {0, 1};
+  const std::vector<std::string> runtime_input_names = {"singleton", "full"};
+
+  DynamicBatchResolutionResult result = ResolveForTest(
+      xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names);
+
+  EXPECT_TRUE(result.can_run);
+  ASSERT_TRUE(result.has_batch_size);
+  EXPECT_EQ(result.batch_size, 240);
+}
+
+TEST(DynamicRuntimeSolveTest, KeepsConsistentNonSingletonEvidence) {
+  const xla::DExpr variable = xla::DExpr::Var(1);
+  const std::vector<xla::Shape> xla_shapes = {
+      MakeDynamicXlaInput(variable),
+      MakeDynamicXlaInput(variable / 3),
+  };
+  const std::vector<TensorShape> runtime_shapes = {
+      TensorShape({720, 16}),
+      TensorShape({240, 16}),
+  };
+  const std::vector<int> runtime_input_indices = {0, 1};
+  const std::vector<std::string> runtime_input_names = {"full", "third"};
+
+  DynamicBatchResolutionResult result = ResolveForTest(
+      xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names);
+
+  EXPECT_TRUE(result.can_run);
+  ASSERT_TRUE(result.has_batch_size);
+  EXPECT_EQ(result.batch_size, 720);
+}
+
+TEST(DynamicRuntimeSolveTest, RejectsConflictingNonSingletonEvidence) {
+  const std::vector<xla::Shape> xla_shapes = {
+      MakeDynamicXlaInput(xla::DExpr::Var(1)),
+      MakeDynamicXlaInput(xla::DExpr::Var(1)),
+  };
+  const std::vector<TensorShape> runtime_shapes = {
+      TensorShape({240, 16}),
+      TensorShape({720, 16}),
+  };
+  const std::vector<int> runtime_input_indices = {0, 1};
+  const std::vector<std::string> runtime_input_names = {"first", "second"};
+
+  DynamicBatchResolutionResult result = ResolveForTest(
+      xla_shapes, runtime_shapes, runtime_input_indices, runtime_input_names);
+
+  EXPECT_FALSE(result.can_run);
+  EXPECT_FALSE(result.has_batch_size);
+  EXPECT_NE(result.diagnostic.find("conflicting non-singleton candidates"),
+            std::string::npos);
+}
+
+TEST(DynamicSolveFilterTest, IgnoresSingletonWhenNonSingletonEvidenceExists) {
+  std::vector<XlaArgument> args = {
+      MakeDynamicArgument(1),
+      MakeDynamicArgument(240),
+  };
+
+  DynamicSolveFilterDecision decision =
+      AnalyzeIgnoredDynamicArgumentOccurrences(args);
+
+  ASSERT_TRUE(decision.can_run);
+  ASSERT_EQ(decision.ignored_occurrences.size(), 1);
+  EXPECT_EQ(decision.ignored_occurrences[0].arg_index, 0);
+  EXPECT_EQ(decision.ignored_occurrences[0].observed_value, 1);
+
+  StripIgnoredDynamicArgumentOccurrences(decision.ignored_occurrences, &args);
+  const TensorShape& singleton_shape = std::get<TensorShape>(args[0].shape);
+  const TensorShape& nonsingleton_shape = std::get<TensorShape>(args[1].shape);
+  EXPECT_FALSE(singleton_shape.get_expression(0));
+  EXPECT_EQ(singleton_shape.dim_size(0), 1);
+  EXPECT_TRUE(nonsingleton_shape.get_expression(0));
+}
+
+TEST(DynamicSolveFilterTest,
+     IgnoresSingletonAcrossExpressionsOfTheSameVariable) {
+  const xla::DExpr variable = xla::DExpr::Var(1);
+  std::vector<XlaArgument> args = {
+      MakeDynamicArgument(1, variable),
+      MakeDynamicArgument(241, variable + 1),
+  };
+
+  DynamicSolveFilterDecision decision =
+      AnalyzeIgnoredDynamicArgumentOccurrences(args);
+
+  ASSERT_TRUE(decision.can_run);
+  ASSERT_EQ(decision.ignored_occurrences.size(), 1);
+  EXPECT_EQ(decision.ignored_occurrences[0].arg_index, 0);
+  EXPECT_EQ(decision.ignored_occurrences[0].observed_value, 1);
+  EXPECT_EQ(decision.ignored_occurrences[0].solved_value, 1);
+}
+
+TEST(DynamicSolveFilterTest, RejectsConflictingNonSingletonEvidence) {
+  std::vector<XlaArgument> args = {
+      MakeDynamicArgument(240),
+      MakeDynamicArgument(720),
+  };
+
+  DynamicSolveFilterDecision decision =
+      AnalyzeIgnoredDynamicArgumentOccurrences(args);
+
+  EXPECT_FALSE(decision.can_run);
+  EXPECT_TRUE(decision.ignored_occurrences.empty());
+  EXPECT_NE(decision.diagnostic.find("conflicting non-singleton candidates"),
+            std::string::npos);
+}
+
+TEST(DynamicSolveFilterTest, KeepsConsistentSingletonOnlyEvidence) {
+  std::vector<XlaArgument> args = {
+      MakeDynamicArgument(1),
+      MakeDynamicArgument(1),
+  };
+
+  DynamicSolveFilterDecision decision =
+      AnalyzeIgnoredDynamicArgumentOccurrences(args);
+
+  EXPECT_TRUE(decision.can_run);
+  EXPECT_TRUE(decision.ignored_occurrences.empty());
+}
+
+TEST(DynamicSolveFilterTest, StaticArgumentsKeepConcreteRuntimeShapes) {
+  XlaArgument dynamic_arg = MakeDynamicArgument(37);
+  dynamic_arg.constant_value_expressions.resize(1);
+  dynamic_arg.constant_value_expressions[0].set_variable_id(1);
+
+  std::vector<XlaArgument> args = {dynamic_arg};
+  std::vector<XlaArgument> static_args =
+      BuildStaticCompilationArguments(args);
+
+  ASSERT_EQ(static_args.size(), 1);
+  const TensorShape& static_shape =
+      std::get<TensorShape>(static_args[0].shape);
+  EXPECT_EQ(static_shape.dim_size(0), 37);
+  EXPECT_EQ(static_shape.dim_size(1), 16);
+  EXPECT_TRUE(static_shape.get_expressions().empty());
+  EXPECT_TRUE(static_args[0].constant_value_expressions.empty());
+}
+
+TEST(DynamicSolveFilterTest, StaticArgumentsClearXlaShapeDynamism) {
+  xla::Shape dynamic_shape;
+  dynamic_shape.set_element_type(xla::F32);
+  dynamic_shape.add_dimensions(37, /*is_dynamic=*/true,
+                               xla::DExpr::Var(1));
+
+  XlaArgument dynamic_arg;
+  dynamic_arg.kind = XlaArgument::kParameter;
+  dynamic_arg.type = DT_FLOAT;
+  dynamic_arg.shape = dynamic_shape;
+  std::vector<XlaArgument> args = {dynamic_arg};
+
+  std::vector<XlaArgument> static_args =
+      BuildStaticCompilationArguments(args);
+
+  const xla::Shape& static_shape = std::get<xla::Shape>(static_args[0].shape);
+  EXPECT_EQ(static_shape.dimensions(0), 37);
+  EXPECT_FALSE(static_shape.is_dynamic_dimension(0));
+  ASSERT_EQ(static_shape.expressions().size(), 1);
+  EXPECT_TRUE(static_shape.expressions(0)->is_constant());
+  EXPECT_EQ(static_shape.expressions(0)->get_val(), 37);
+}
+
+}  // namespace
+}  // namespace xla_ops_internal
+}  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc
index c3a24f3e0f7163..d1343c7235fc22 100644
--- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc
+++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc
@@ -23,12 +23,14 @@ limitations under the License.
 #include <map>
 #include <memory>
 #include <optional>
+#include <regex>
 #include <set>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
 #include <utility>
 #include <vector>
+#include <queue>
 
 #include "absl/base/call_once.h"
 #include "absl/container/flat_hash_map.h"
@@ -39,6 +41,7 @@ limitations under the License.
 #include "tensorflow/compiler/jit/deadness_analysis.h"
 #include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/jit/device_util.h"
+#include "tensorflow/compiler/jit/encapsulate_util.h"
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/resource_operation_safety_analysis.h"
 #include "tensorflow/compiler/jit/xla_cluster_util.h"
@@ -68,6 +71,9 @@ limitations under the License.
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/public/version.h"
 #include "tensorflow/core/util/dump_graph.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/grappler/costs/graph_properties.h"
+#include "tensorflow/core/grappler/grappler_item.h"
 
 namespace tensorflow {
 
@@ -109,6 +115,8 @@ class MarkForCompilationPassImpl {
     // stable from run to rum.
     bool deterministic_cluster_names;
 
+    bool enable_dynamic_sizes;
+
     int max_cluster_size;
     int min_cluster_size;
 
@@ -123,6 +131,11 @@ class MarkForCompilationPassImpl {
     std::atomic<int64_t>* fuel;
 
     bool dump_graphs;
+
+    // Enable models to influcence clustering with operator names
+    int annotate_cluster_id;
+
+    bool enable_cluster_parallel;
   };
 
   MarkForCompilationPassImpl(DebugOptions debug_options, Graph* graph,
@@ -245,7 +258,28 @@ class MarkForCompilationPassImpl {
                           " others #", cycles_graph_node_id(), ">");
     }
 
+    int annotated_id() const { return annotated_id_; }
+    void set_annotated_id(int id) { annotated_id_ = id; }
+    int chain_id() const {return chain_id_;}
+    void set_chain_id(int id) {chain_id_ = id;}
+    void add_dim_var(int dim_var) { dim_vars_.insert(dim_var); }
+    void merge_dim_vars(const std::set<int>& dim_vars) {
+      dim_vars_.insert(dim_vars.begin(), dim_vars.end());
+    }
+    const std::set<int>& dim_vars() const { return dim_vars_; }
+    void add_dim_expr(const xla::DExpr& dim_expr) {
+      dim_exprs_.push_back(dim_expr);
+    }
+    void merge_dim_exprs(const std::vector<xla::DExpr>& dim_exprs) {
+      dim_exprs_.insert(dim_exprs_.end(), dim_exprs.begin(), dim_exprs.end());
+    }
+    const std::vector<xla::DExpr>& dim_exprs() const { return dim_exprs_; }
+
    private:
+    int annotated_id_ = -1;
+    std::set<int> dim_vars_;
+    std::vector<xla::DExpr> dim_exprs_;
+    int chain_id_ = -1;
     int cluster_size_ = 1;
     int cycles_graph_node_id_;
     int effective_cluster_size_;
@@ -317,6 +351,21 @@ class MarkForCompilationPassImpl {
     return compilation_candidates_.find(n) != compilation_candidates_.end();
   }
 
+  absl::Status AssignAnnotatedClusterIDs();
+  absl::Status AssignDimVars();
+  std::optional<std::string> CheckDynamicExpressionCompatibility(
+      absl::Span<const xla::DExpr> exprs);
+  bool DynamicNodeExpressionsAreCompatible(const Node& node,
+                                           std::string* reason);
+  void collectInputNodes(std::set<Node*> &path_nodes);
+  void collectMergeNodes(const std::vector<Node*>& nodeSet,
+                         std::set<Node*> &merger_nodes);
+  void collectPathNodes(Node* start, std::set<Node*> &path_nodes,
+                        std::set<Node*>& merger_nodes);
+  std::map<Node*, std::vector<Node*>> collectParallelNode(
+      const std::vector<Node*>& nodeSet);
+  absl::Status AssignParallelChains();
+
   // Tries to contract the edge from cluster `from` to cluster `to`.  Returns
   // true if successful.
   absl::StatusOr<bool> TryToContractEdge(Cluster* from, Cluster* to);
@@ -606,6 +655,11 @@ void MarkForCompilationPassImpl::Cluster::Merge(Cluster* other) {
     xla_scope_ = std::move(other->xla_scope_);
   }
 
+  merge_dim_vars(other->dim_vars_);
+  other->dim_vars_.clear();
+  merge_dim_exprs(other->dim_exprs_);
+  other->dim_exprs_.clear();
+
   resource_var_operation_node_ids_.reserve(
       resource_var_operation_node_ids_.size() +
       other->resource_var_operation_node_ids_.size());
@@ -651,11 +705,320 @@ absl::Status IgnoreResourceOpForSafetyAnalysis(
   }
   return absl::OkStatus();
 }
+// node mapping to multiple vectors of expressions (one for each output in
+// order)
+static std::map<std::string, std::vector<std::vector<std::unique_ptr<DimExpr>>>>
+    expr_map;
+// Helper to convert ExpressionProto to a readable string.
+std::string ExprProtoToString(const ExpressionProto& e) {
+  switch (e.node_type_case()) {
+    case ExpressionProto::kConstantValue:
+      return std::to_string(e.constant_value());
+    case ExpressionProto::kVariableId:
+      return absl::StrCat("Var(", e.variable_id(), ")");
+    case ExpressionProto::kAddNode:
+      return absl::StrCat("(", ExprProtoToString(e.add_node().lhs()), " + ",
+                         ExprProtoToString(e.add_node().rhs()), ")");
+    case ExpressionProto::kSubNode:
+      return absl::StrCat("(", ExprProtoToString(e.sub_node().lhs()), " - ",
+                         ExprProtoToString(e.sub_node().rhs()), ")");
+    case ExpressionProto::kMulNode:
+      return absl::StrCat("(", ExprProtoToString(e.mul_node().lhs()), " * ",
+                         ExprProtoToString(e.mul_node().rhs()), ")");
+    case ExpressionProto::kDivNode:
+      return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ",
+                         ExprProtoToString(e.div_node().rhs()), ")");
+    case ExpressionProto::kMaxNode:
+      return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ",
+                          ExprProtoToString(e.max_node().rhs()), ")");
+    case ExpressionProto::kGtNode:
+      return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ",
+                          ExprProtoToString(e.gt_node().rhs()), ")");
+    case ExpressionProto::kSelectNode:
+      return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()),
+                          ", ", ExprProtoToString(e.select_node().on_true()),
+                          ", ", ExprProtoToString(e.select_node().on_false()),
+                          ")");
+    default:
+      return "<none>";
+  }
+}
+
+bool HasDynamicInputExpression(const Node* node) {
+  for (const Edge* edge : node->in_edges()) {
+    if (edge->IsControlEdge()) {
+      continue;
+    }
+    auto it = expr_map.find(edge->src()->name());
+    if (it == expr_map.end()) {
+      continue;
+    }
+    const int output_index = edge->src_output();
+    if (output_index < 0 || output_index >= it->second.size()) {
+      continue;
+    }
+    for (const auto& expr : it->second[output_index]) {
+      if (expr == nullptr) {
+        continue;
+      }
+      xla::DExpr dynamic_expr = *expr;
+      if (dynamic_expr && dynamic_expr->is_dynamic()) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+std::string DynamicExpressionToString(const xla::DExpr& expr) {
+  xla::DExpr simplified = expr.simplify();
+  if (!simplified && !simplified.is_unknown()) {
+    return "";
+  }
+  xla::StringPrinter printer;
+  simplified->print(&printer);
+  return std::move(printer).ToString();
+}
+
+std::optional<std::string> CheckDynamicExpressionCompatibilityImpl(
+    absl::Span<const xla::DExpr> exprs) {
+  const xla::DExpr* anchor_source = nullptr;
+  for (const xla::DExpr& expr : exprs) {
+    if (expr && expr->is_dynamic() && !expr->get_all_ids().empty()) {
+      anchor_source = &expr;
+      break;
+    }
+  }
+  if (anchor_source == nullptr) {
+    return std::nullopt;
+  }
+
+  const std::set<int> expected_ids = (*anchor_source)->get_all_ids();
+  std::optional<xla::DExpr> expected_core;
+  if (expected_ids.size() > 1) {
+    expected_core =
+        anchor_source->find_smallest_subexpression_covering_all_variables();
+  }
+  int replacement_id = std::numeric_limits<int>::min();
+  while (expected_ids.count(replacement_id) != 0) {
+    ++replacement_id;
+  }
+  const std::set<int> replacement_ids = {replacement_id};
+
+  for (const xla::DExpr& expr : exprs) {
+    if (!expr || !expr->is_dynamic() || expr->get_all_ids().empty()) {
+      continue;
+    }
+    const std::set<int> expr_ids = expr->get_all_ids();
+    if (expr_ids != expected_ids) {
+      return absl::StrCat(
+          "dynamic expressions do not share the same variable ids: "
+          "expected={",
+          absl::StrJoin(expected_ids, ", "), "}, expr=",
+          DynamicExpressionToString(expr), ", ids={",
+          absl::StrJoin(expr_ids, ", "), "}");
+    }
+    if (!expected_core.has_value()) {
+      continue;
+    }
+    const xla::DExpr core =
+        expr.find_smallest_subexpression_covering_all_variables();
+    if (!(core == *expected_core)) {
+      return absl::StrCat(
+          "dynamic expressions do not share the same clusterable core: "
+          "expected=",
+          DynamicExpressionToString(*expected_core), ", expr=",
+          DynamicExpressionToString(expr), ", core=",
+          DynamicExpressionToString(core));
+    }
+    const xla::DExpr normalized =
+        expr.replace_subexpression(core, xla::DExpr::Var(replacement_id))
+            .simplify();
+    if (normalized->get_all_ids() != replacement_ids) {
+      return absl::StrCat(
+          "dynamic expression core does not cover every variable occurrence: "
+          "expr=",
+          DynamicExpressionToString(expr), ", core=",
+          DynamicExpressionToString(core));
+    }
+  }
+
+  return std::nullopt;
+}
+
+std::optional<std::string>
+MarkForCompilationPassImpl::CheckDynamicExpressionCompatibility(
+    absl::Span<const xla::DExpr> exprs) {
+  return CheckDynamicExpressionCompatibilityImpl(exprs);
+}
+
+bool MarkForCompilationPassImpl::DynamicNodeExpressionsAreCompatible(
+    const Node& node, std::string* reason) {
+  std::vector<xla::DExpr> node_exprs;
+
+  for (const Edge* edge : node.in_edges()) {
+    if (edge->IsControlEdge()) {
+      continue;
+    }
+
+    const Node* src = edge->src();
+    auto it = expr_map.find(src->name());
+    if (it == expr_map.end()) {
+      continue;
+    }
+
+    const int output_index = edge->src_output();
+    if (output_index < 0 ||
+        output_index >= static_cast<int>(it->second.size())) {
+      continue;
+    }
+
+    for (const auto& expr_ptr : it->second[output_index]) {
+      if (expr_ptr == nullptr) {
+        continue;
+      }
+      xla::DExpr dyn = *expr_ptr;
+      if (!dyn || !dyn->is_dynamic() || dyn->get_all_ids().empty()) {
+        continue;
+      }
+      node_exprs.push_back(std::move(dyn));
+    }
+  }
+
+  auto output_it = expr_map.find(node.name());
+  if (output_it != expr_map.end()) {
+    for (const auto& output_exprs : output_it->second) {
+      for (const auto& expr_ptr : output_exprs) {
+        if (expr_ptr == nullptr) {
+          continue;
+        }
+        xla::DExpr dyn = *expr_ptr;
+        if (!dyn || !dyn->is_dynamic() || dyn->get_all_ids().empty()) {
+          continue;
+        }
+        node_exprs.push_back(std::move(dyn));
+      }
+    }
+  }
+
+  const std::optional<std::string> incompatibility =
+      CheckDynamicExpressionCompatibility(node_exprs);
+  if (!incompatibility.has_value()) {
+    return true;
+  }
+  *reason = *incompatibility;
+  return false;
+}
+
+// Runs Grappler static inference and logs any ExpressionProto found in output
+// tensor shapes (from GraphProperties, not from _output_shapes attrs).
+void LogExpressionsViaGraphProperties(tensorflow::Graph& graph) {
+  using tensorflow::ExpressionProto;
+  using tensorflow::GraphDef;
+  using tensorflow::NodeDef;
+  using tensorflow::TensorShapeProto;
+  using tensorflow::grappler::GraphProperties;
+  using tensorflow::grappler::GrapplerItem;
+
+  expr_map.clear();
+
+  GraphDef graph_def;
+  graph.ToGraphDef(&graph_def);
+  auto node_name_index = graph.BuildNodeNameIndex();
+
+  GrapplerItem item;
+  item.id = "mark_for_compilation_pass_expr_dump";
+  item.graph = graph_def;
+
+  GraphProperties props(item);
+
+  absl::Status st = props.InferStatically(
+      /*assume_valid_feeds=*/false,
+      /*aggressive_shape_inference=*/false,
+      /*include_input_tensor_values=*/false,
+      /*include_output_tensor_values=*/false,
+      /*enable_dynamic_value_inference=*/true);
+
+  if (!st.ok()) {
+    LOG(ERROR) << "[EXPR][GP] InferStatically failed: " << st.message();
+    return;
+  }
+
+  int found = 0;
+  VLOG(1) << "[EXPR][GP] === GraphProperties output expr dump ===";
+
+  auto convert_graph_properties_shape = [](const TensorShapeProto& gp_shape) {
+    TensorShapeProto out;
+    out.set_unknown_rank(gp_shape.unknown_rank());
+    for (int i = 0; i < gp_shape.dim_size(); ++i) {
+      const auto& dim = gp_shape.dim(i);
+      out.add_dim()->set_size(dim.size());
+      ExpressionProto* expr = out.add_expressions();
+      if (i < gp_shape.expressions_size() &&
+          gp_shape.expressions(i).node_type_case() !=
+              ExpressionProto::NODE_TYPE_NOT_SET) {
+        *expr = gp_shape.expressions(i);
+      } else {
+        expr->set_constant_value(dim.size());
+      }
+    }
+    return out;
+  };
+
+  for (const NodeDef& n : graph_def.node()) {
+    if (!props.HasOutputProperties(n.name())) continue;
+    const auto& outs = props.GetOutputProperties(n.name());
+    std::vector<TensorShapeProto> inferred_output_shapes;
+    inferred_output_shapes.reserve(outs.size());
+    std::vector<std::vector<std::unique_ptr<DimExpr>>> list_exprs(outs.size());
+    for (int out_idx = 0; out_idx < static_cast<int>(outs.size()); ++out_idx) {
+      const auto& tp = outs[out_idx];
+      const TensorShapeProto& shp = tp.shape();
+      inferred_output_shapes.push_back(convert_graph_properties_shape(shp));
+
+      std::vector<std::unique_ptr<DimExpr>> exprs;
+      for (int d = 0; d < shp.dim_size(); ++d) {
+        if (d >= shp.expressions_size()) continue;
+        const ExpressionProto& expr = shp.expressions(d);
+        if (expr.node_type_case() == ExpressionProto::NODE_TYPE_NOT_SET)
+          continue;
+
+        VLOG(1) << "Node " << n.name() << " has expression "
+                << ExprProtoToString(expr);
+
+        exprs.push_back(
+            std::make_unique<DimExpr>(DimExprFromProto(expr)));
+
+        ++found;
+      }
+      if (shp.dim_size() == 0 && shp.unknown_rank()) {
+        // Add two dummy variables to represent the unknown rank
+        exprs.push_back(std::make_unique<DimExpr>(DimExpr::Var(-888)));
+        exprs.push_back(std::make_unique<DimExpr>(DimExpr::Var(-889)));
+      }
+
+      list_exprs[out_idx] = std::move(exprs);
+    }
+    expr_map[n.name()] = std::move(list_exprs);
+    auto node_it = node_name_index.find(n.name());
+    if (node_it != node_name_index.end()) {
+      node_it->second->AddAttr(kXlaInferredOutputTensorShapesAttrName,
+                               inferred_output_shapes);
+    }
+
+  }
+
+  VLOG(1) << "[EXPR][GP] === Found " << found
+          << " expressions via GraphProperties ===";
+}
 
 absl::StatusOr<bool> MarkForCompilationPassImpl::Initialize() {
   TF_RET_CHECK(!initialized_ && !edges_contracted_ && !clusters_created_);
   initialized_ = true;
 
+  if (debug_options_.enable_dynamic_sizes) {
+    LogExpressionsViaGraphProperties(*graph_);
+  }
   TF_RETURN_IF_ERROR(FindCompilationCandidates());
 
   if (compilation_candidates_.empty()) {
@@ -686,6 +1049,44 @@ absl::StatusOr<bool> MarkForCompilationPassImpl::Initialize() {
   // representative names the node in the 'cycles' graph that represents the
   // cluster.
   TF_RETURN_IF_ERROR(BuildInitialClusterSet());
+
+  // Source model may be annotated with preferred clusters. This function
+  // just interpreter the annotations and assign preferred IDs
+  if (debug_options_.annotate_cluster_id) {
+    TF_RETURN_IF_ERROR(AssignAnnotatedClusterIDs());
+  }
+  if (debug_options_.enable_dynamic_sizes) {
+    TF_RETURN_IF_ERROR(AssignDimVars());
+    for (Node* n : graph_->op_nodes()) {
+      bool mark_shape_derived = false;
+      auto is_shape_like = [](const Node* node) {
+        const string& type = node->type_string();
+        return type == "Shape" || type == "ShapeN" || type == "Size";
+      };
+      if (is_shape_like(n)) {
+        mark_shape_derived = HasDynamicInputExpression(n);
+      } else if (n->type_string() == "Cast") {
+        for (const Edge* edge : n->in_edges()) {
+          if (edge->IsControlEdge()) {
+            continue;
+          }
+          const Node* src = edge->src();
+          if (is_shape_like(src) && HasDynamicInputExpression(src)) {
+            mark_shape_derived = true;
+            break;
+          }
+        }
+      }
+      if (mark_shape_derived) {
+        n->AddAttr(kXlaShapeDerivedAttrName, true);
+        VLOG(1) << "MarkForCompilation marked shape-derived node "
+                << n->name() << " op=" << n->type_string();
+      }
+    }
+  }
+  if (debug_options_.enable_cluster_parallel) {
+    TF_RETURN_IF_ERROR(AssignParallelChains());
+  }
   return true;
 }
 
@@ -893,7 +1294,12 @@ absl::Status MarkForCompilationPassImpl::RunEdgeContractionLoop() {
                       ForEachEdgeInPostOrder([&](Cluster* from, Cluster* to) {
                         return TryToContractEdge(from, to);
                       }));
+  /* Clustering conditions for dynamic shapes may break the assumption of
+   * fixed point at phase 2, so this check may fail.
+   * Now just disable this check, and we can re-enable it after we have more
+   * confidence in the code.
   TF_RET_CHECK(!changed);
+   */
 
   return absl::OkStatus();
 }
@@ -1028,6 +1434,7 @@ absl::Status MarkForCompilationPassImpl::CreateClusters() {
     // trouble.
 
     if (cluster->effective_cluster_size() >= debug_options_.min_cluster_size ||
+        cluster->chain_id() != -1 ||
         cluster->has_functional_control_flow() ||
         cluster->is_xla_compile_attr_true()) {
       string& name = cluster_names[cluster->cycles_graph_node_id()];
@@ -1414,6 +1821,15 @@ absl::Status MarkForCompilationPassImpl::FindCompilationCandidates() {
       continue;
     }
 
+    if (debug_options_.enable_dynamic_sizes &&
+        XlaOpRegistry::IsMlirXlaOp(node->type_string()) &&
+        HasDynamicInputExpression(node)) {
+      VLOG(1) << "Rejecting " << node->name()
+              << " from XLA clustering: MlirXlaOpKernel does not support "
+                 "dynamic input expressions";
+      continue;
+    }
+
     if (node->type_string() == "Const") {
       // Skip Const op with type DT_STRING, since XLA autoclustering doesn't
       // support it.
@@ -1429,6 +1845,15 @@ absl::Status MarkForCompilationPassImpl::FindCompilationCandidates() {
       continue;
     }
 
+    if (debug_options_.enable_dynamic_sizes) {
+      std::string reason;
+      if (!DynamicNodeExpressionsAreCompatible(*node, &reason)) {
+        VLOG(1) << "Rejecting " << node->name()
+                << " from XLA clustering: " << reason;
+        continue;
+      }
+    }
+
     if (compile_time_const_nodes[node->id()]) {
       const OpDef* op_def;
       TF_RETURN_IF_ERROR(
@@ -1548,14 +1973,322 @@ bool MarkForCompilationPassImpl::CompilationDisallowedByXlaCompileAttr(
   return false;
 }
 
+absl::Status MarkForCompilationPassImpl::AssignAnnotatedClusterIDs(void) {
+  VLOG(1) << "Run AssignAnnotatedClusterIDs";
+
+  for (Node* node : graph_->nodes()) {
+    Cluster * cluster = GetClusterForNode(node);
+    auto name = node->name();
+    if (cluster) {
+      std::string pat = "^\\.cluster\\.(\\d+|none)";
+      std::regex idPattern(pat);
+      std::smatch matched;
+      if (std::regex_search(name, matched, idPattern)) {
+	auto m = matched.str(1);
+	if (m == "none") {
+          // Prefer not to cluster
+	  VLOG(1) << name << " : Default annotated cluster id -1";
+	  cluster->set_annotated_id(-1);
+	}
+	else {
+	  try {
+	    int id = std::stoi(m);
+	    cluster->set_annotated_id(id);
+	    VLOG(1) << name << " : Set annotated cluster id " << m;
+	  }
+	  catch (...) {
+            VLOG(1) << name << " : Invalid cluster id: " << m;
+	  }
+	}
+      }
+      else {
+        VLOG(1) << "Not matched: " << name << " pattern is " << pat;
+      }
+    }
+    else {
+      VLOG(1) << name << ": Not initially clustered";
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status MarkForCompilationPassImpl::AssignDimVars(void) {
+  for (Node* node : graph_->nodes()) {
+    auto node_name = node->name();
+    Cluster * cluster = GetClusterForNode(node);
+    if (!cluster) continue;
+    for (const tensorflow::Edge* edge : node->in_edges()) {
+      if (edge->IsControlEdge()) {
+          // Skip control edges if you are only interested in data edges
+          continue;
+      }
+
+      const tensorflow::Node* input = edge->src(); // Source node of the edge
+      auto it = expr_map.find(input->name());
+      if (it == expr_map.end()) {
+        VLOG(2) << "No expression found for node " << input->name();
+        continue;
+      }
+
+      auto output_index = edge->src_output(); // Output index of the source node
+      if (output_index >= (it->second).size()) {
+        LOG(INFO) << "Warning: Output index " << output_index << " is out of bounds for node " << input->name();
+        continue;
+      }
+      for (auto& pDim: (it->second)[output_index]) {
+        DimExpr * d= pDim.get();
+        xla::DExpr dyn = *d;
+        if (!dyn || !dyn->is_dynamic()) {
+          continue;
+        }
+        cluster->add_dim_expr(dyn);
+        auto new_ids = dyn->get_all_ids();
+        for (auto id : new_ids) {
+          cluster->add_dim_var(id);
+          VLOG(2) << "Add dim var " << id << " to cluster of node "<< node_name;
+        }
+      }
+    }
+    auto output_it = expr_map.find(node_name);
+    if (output_it != expr_map.end()) {
+      for (const auto& output_exprs : output_it->second) {
+        for (const auto& expr_ptr : output_exprs) {
+          if (expr_ptr == nullptr) {
+            continue;
+          }
+          xla::DExpr dyn = *expr_ptr;
+          if (!dyn || !dyn->is_dynamic()) {
+            continue;
+          }
+          cluster->add_dim_expr(dyn);
+        }
+      }
+    }
+    // create a for loop for each dim vars in cluster and print each dim var
+    if (VLOG_IS_ON(2)) {
+      if (cluster->dim_vars().empty()) {
+        VLOG(2) << "Cluster of node " << node_name << " has no dim vars.";
+      }
+      else {
+        std::string id_str;
+        for (auto id : cluster->dim_vars()) {
+          id_str += "Dim var " + std::to_string(id) + ", ";
+        }
+        VLOG(2) << "Cluster of node " << node_name << " has dim vars:\n" << id_str;
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
 bool MarkForCompilationPassImpl::LogNotContractableAndReturnFalse(
     Cluster* from, Cluster* to, absl::string_view reason) {
   VLOG(3) << EdgeContractionFailureMsg(from, to, reason);
   return false;
 }
 
+void MarkForCompilationPassImpl::collectInputNodes(std::set<Node*> &path_nodes) {
+  std::unordered_map<Node*, int> out_degree_count;
+
+  // 4. Initialize the queue and add nodes from path_nodes
+  std::queue<Node*> queue;
+  for (auto node : path_nodes) {
+    queue.push(node);
+  }
+
+  // 5. BFS search
+  while (!queue.empty()) {
+    auto u = queue.front();
+    queue.pop();
+
+    // Traverse all predecessor nodes
+    for (const Edge* e : u->in_edges()) {
+      Node* p = e->src();
+      if (path_nodes.find(p) != path_nodes.end() ||
+          !IsCompilationCandidate(p)) {
+        continue;
+      }
+      if (out_degree_count.count(p) == 0) {
+        // Initialize the out-degree count for the node
+        out_degree_count[p] = p->out_edges().size();
+      }
+      // Decrease the out-degree count for the predecessor node
+      out_degree_count[p] -= 1;
+
+      // If the predecessor node's out-degree count is 0 and not in path_nodes,
+      // add it to path_nodes and queue
+      if (out_degree_count[p] == 0) {
+        path_nodes.insert(p);
+        queue.push(p);
+        VLOG(3) << p->DebugString();
+      }
+    }
+  }
+}
+
+void MarkForCompilationPassImpl::collectMergeNodes(
+    const std::vector<Node*>& nodeSet, std::set<Node*> &merger_nodes) {
+  // 1. Collect the number of nodeSet that can reach each node
+  std::map<Node*, std::set<Node*>> reach_map;
+  for (Node* start : nodeSet) {
+    std::set<Node*> visited;
+    std::vector<Node*> stack = {start};
+    while (!stack.empty()) {
+      Node* cur = stack.back();
+      stack.pop_back();
+      if (visited.count(cur)) continue;
+      visited.insert(cur);
+      for (const Edge* e : cur->out_edges()) {
+        Node* next = e->dst();
+        if (!visited.count(next))
+          stack.push_back(next);
+      }
+    }
+    reach_map[start] = std::move(visited);
+  }
+
+  // 2. Determine the merger node
+  std::map<Node*, int> node_reach_count;
+  std::set<Node*> all_nodes;
+  for (const auto& kv : reach_map) {
+    for (Node* n : kv.second) {
+      node_reach_count[n]++;
+      all_nodes.insert(n);
+    }
+  }
+  for (Node* n : all_nodes) {
+    // Condition 1: Reached by multiple sources
+    if (node_reach_count[n] >= 2) merger_nodes.insert(n);
+    // Condition 2: No output edges
+    if (n->out_edges().empty()) merger_nodes.insert(n);
+  }
+}
+
+void MarkForCompilationPassImpl::collectPathNodes(
+    Node* start, std::set<Node*> &path_nodes, std::set<Node*>& merger_nodes) {
+  std::vector<Node*> stack = {start};
+  std::set<Node*> visited;
+
+  while (!stack.empty()) {
+    Node* cur = stack.back();
+    stack.pop_back();
+    if (visited.count(cur) || !IsCompilationCandidate(cur)) {
+      continue;
+    }
+    visited.insert(cur);
+
+    // Stop search met the merger node
+    if (merger_nodes.count(cur)) {
+      if (cur->out_edges().empty()) path_nodes.insert(cur);
+      continue;
+    }
+    path_nodes.insert(cur);
+
+    for (const Edge* e : cur->out_edges()) {
+      Node* next = e->dst();
+      if (!visited.count(next)) {
+        stack.push_back(next);
+      }
+    }
+  }
+}
+
+// collectParallelNode
+// Search the serial merger nodes based on the parallel matmul starting points
+// Search along the output edge to get the boundary from start to all merger points
+// Search along the input edge to get the entire parallel computation graph
+std::map<Node*, std::vector<Node*>>
+MarkForCompilationPassImpl::collectParallelNode(
+    const std::vector<Node*>& nodeSet) {
+  std::set<Node*> merger_nodes;
+  collectMergeNodes(nodeSet, merger_nodes);
+
+  // Collect path nodes
+  std::map<Node*, std::vector<Node*>> result;
+  for (Node* start : nodeSet) {
+    VLOG(4) << "Search parallel graph form node: " << start->DebugString();
+    std::set<Node*> path_nodes;
+
+    // Search along output edge form start to merger nodes
+    collectPathNodes(start, path_nodes, merger_nodes);
+
+    VLOG(4) << "Collect path nodes:";
+    for (auto node : path_nodes) {
+      VLOG(4) << node->type_string();
+    }
+
+    VLOG(4) << "Collect input nodes:";
+    // search along input edge
+    collectInputNodes(path_nodes);
+
+    result[start] = std::vector<Node*>(path_nodes.begin(), path_nodes.end());
+  }
+  return result;
+}
+
+// Collect parallel matmuls as input nodes into nodeSet
+// Use collectParallelNode to get the parallel subgraph
+// Mark parallel nodes change ID
+absl::Status MarkForCompilationPassImpl::AssignParallelChains() {
+  VLOG(4) << "Run AssignParallelChains";
+  // Record the matmuls that can be paralleled
+  std::vector<std::vector<Node*>> parallel_matmuls;
+  int minParallelMatmulNum = 2;
+  int next_chain_id = 0;
+  // Collect matmul nodes with shared input to parallel_matmuls
+  for (Node* node : graph_->nodes()) {
+    if (node->out_edges().size() < 2) continue;
+    std::vector<Node*> matmul_nodes;
+    for (const Edge* e : node->out_edges()) {
+      if (e->IsControlEdge()) continue;
+      Node* succ = e->dst();
+      VLOG(4) << "Find matmul node: " << succ->type_string() << " : " << succ->DebugString();
+      if (succ->type_string() == "MatMul")
+        matmul_nodes.push_back(succ);
+    }
+    if (matmul_nodes.size() >= minParallelMatmulNum)
+      parallel_matmuls.push_back(matmul_nodes);
+  }
+
+  for (auto matmul_nodes : parallel_matmuls) {
+    VLOG(4) << "Process matmul nodes: Total " << matmul_nodes.size()
+            << " sub matmuls";
+    bool visited = false;
+    for (auto matmul : matmul_nodes) {
+      VLOG(4) << matmul->name();
+      Cluster* cluster = GetClusterForNode(matmul);
+      if (!cluster || cluster->chain_id() != -1) {
+        visited = true;
+        break;
+      }
+    }
+    if (visited) {
+      VLOG(4) << "stop collect: has visited matmul";
+      continue;
+    }
+
+    std::map<Node*, std::vector<Node*>> subgraphMap =
+        collectParallelNode(matmul_nodes);
+    for (auto it : subgraphMap) {
+      auto nodeSet = it.second;
+      VLOG(4) << "One of Parallel sub-graph is: ";
+      for (auto node : nodeSet) {
+        VLOG(4) << node->DebugString();
+        Cluster* cluster = GetClusterForNode(node);
+        cluster->set_chain_id(next_chain_id);
+      }
+      next_chain_id++;
+    }
+  }
+  return absl::OkStatus();
+}
+
 absl::StatusOr<bool> MarkForCompilationPassImpl::TryToContractEdge(
     Cluster* from, Cluster* to) {
+  if (from->chain_id() != to->chain_id()) {
+    return LogNotContractableAndReturnFalse(
+        from, to, "nodes are in different parallel chains");
+  }
   DCHECK(from->deadness_predicate().has_value() ==
          to->deadness_predicate().has_value());
   if (from->deadness_predicate() != to->deadness_predicate()) {
@@ -1569,6 +2302,22 @@ absl::StatusOr<bool> MarkForCompilationPassImpl::TryToContractEdge(
     return false;
   }
 
+  if (debug_options_.annotate_cluster_id && from->annotated_id() != to->annotated_id()) {
+    return LogNotContractableAndReturnFalse(
+        from, to, "the two nodes do not have same annotated ids");
+  }
+
+  if (debug_options_.enable_dynamic_sizes) {
+    std::vector<xla::DExpr> combined_exprs = from->dim_exprs();
+    combined_exprs.insert(combined_exprs.end(), to->dim_exprs().begin(),
+                          to->dim_exprs().end());
+    std::optional<std::string> incompatibility =
+        CheckDynamicExpressionCompatibility(combined_exprs);
+    if (incompatibility.has_value()) {
+      return LogNotContractableAndReturnFalse(from, to, *incompatibility);
+    }
+  }
+
   TF_ASSIGN_OR_RETURN(bool devices_compatible,
                       AreDevicesCompatible(*from, *to));
   if (!devices_compatible) {
@@ -1639,7 +2388,6 @@ absl::Status MarkForCompilationPassImpl::Run() {
     // MarkForCompilationPassImpl is not set up to run the subsequent phases.
     return absl::OkStatus();
   }
-
   TF_RETURN_IF_ERROR(RunEdgeContractionLoop());
   TF_RETURN_IF_ERROR(DeclusterNodes());
   TF_RETURN_IF_ERROR(CreateClusters());
@@ -1958,10 +2706,14 @@ absl::Status MarkForCompilationPass::Run(
   debug_options.ignore_xla_compile_attr = false;
   debug_options.deterministic_cluster_names =
       flags->tf_xla_deterministic_cluster_names;
+  debug_options.enable_dynamic_sizes =
+      flags->tf_xla_enable_dynamic_sizes;
   debug_options.max_cluster_size = flags->tf_xla_max_cluster_size;
   debug_options.min_cluster_size = flags->tf_xla_min_cluster_size;
   debug_options.fuel = GetPointerToFuel(flags->tf_xla_clustering_fuel);
   debug_options.dump_graphs = flags->tf_xla_clustering_debug;
+  debug_options.annotate_cluster_id = flags->tf_xla_annotate_cluster_id;
+  debug_options.enable_cluster_parallel = flags->tf_xla_cluster_parallel;
 
   return MarkForCompilation(options, debug_options);
 }
@@ -1977,10 +2729,12 @@ absl::Status MarkForCompilationPass::RunForTest(
       flags->tf_xla_disable_resource_variable_safety_checks_for_debugging;
   debug_options.ignore_xla_compile_attr = true;
   debug_options.deterministic_cluster_names = deterministic_cluster_names;
+  debug_options.enable_dynamic_sizes = false;
   debug_options.max_cluster_size = flags->tf_xla_max_cluster_size;
   debug_options.min_cluster_size = flags->tf_xla_min_cluster_size;
   debug_options.fuel = GetPointerToFuel(flags->tf_xla_clustering_fuel);
   debug_options.dump_graphs = flags->tf_xla_clustering_debug;
+  debug_options.annotate_cluster_id = flags->tf_xla_annotate_cluster_id;
 
   return MarkForCompilation(options, debug_options);
 }
@@ -2056,6 +2810,11 @@ void ResetClusterSequenceNumber() {
   ClusterSequenceNumberGenerator::Global().Reset();
 }
 
+std::optional<std::string> CheckDynamicExpressionCompatibilityForTest(
+    absl::Span<const xla::DExpr> exprs) {
+  return CheckDynamicExpressionCompatibilityImpl(exprs);
+}
+
 absl::flat_hash_set<string> GetKnownXLAAllowlistOp() {
   absl::flat_hash_set<string> result{
       "AdjustContrastv2",
diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.h b/tensorflow/compiler/jit/mark_for_compilation_pass.h
index 558912f2eee2e0..79e3c88074fcdf 100644
--- a/tensorflow/compiler/jit/mark_for_compilation_pass.h
+++ b/tensorflow/compiler/jit/mark_for_compilation_pass.h
@@ -20,10 +20,17 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_
 #define TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_
 
+#include <optional>
+
 #include "absl/container/flat_hash_set.h"
+#include "absl/types/span.h"
 #include "tensorflow/compiler/jit/compilability_check_util.h"
 #include "tensorflow/core/common_runtime/optimization_registry.h"
 
+namespace xla {
+class DExpr;
+}
+
 namespace tensorflow {
 
 // The attribute that marks nodes to be grouped into functions by the
@@ -57,6 +64,11 @@ void ResetClusterSequenceNumber();
 
 // Return a list of operation that we choose not to put into the allowlist.
 absl::flat_hash_set<string> GetKnownXLAAllowlistOp();
+
+// Returns an explanation when dynamic expressions cannot use the same
+// cluster-level dynamic value.
+std::optional<std::string> CheckDynamicExpressionCompatibilityForTest(
+    absl::Span<const xla::DExpr> exprs);
 }  // namespace testing
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc
index 1a120791206369..bdc4344af22547 100644
--- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc
+++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc
@@ -53,6 +53,7 @@ limitations under the License.
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/test.h"
+#include "xla/shape_expr.h"
 
 using ::tensorflow::testing::FindNodeByName;
 
@@ -117,6 +118,74 @@ absl::flat_hash_map<string, std::vector<string>> GetClusterSets(
   return cluster_sets;
 }
 
+// Expressions may differ outside the smallest subtree containing all variables.
+TEST(XlaCompilationTest,
+     DynamicExpressionCompatibilityUsesSmallestCoveringSubexpression) {
+  const xla::DExpr shared_core =
+      xla::DExpr::Var(1) + xla::DExpr::Var(2);
+  const std::vector<xla::DExpr> compatible_exprs = {
+      shared_core + 2, shared_core + 3};
+
+  EXPECT_FALSE(
+      testing::CheckDynamicExpressionCompatibilityForTest(compatible_exprs)
+          .has_value());
+
+  const std::vector<xla::DExpr> incompatible_exprs = {
+      shared_core + 2,
+      (xla::DExpr::Var(1) - xla::DExpr::Var(2)) + 3};
+  const std::optional<std::string> incompatibility =
+      testing::CheckDynamicExpressionCompatibilityForTest(incompatible_exprs);
+
+  ASSERT_TRUE(incompatibility.has_value());
+  EXPECT_NE(incompatibility->find("same clusterable core"),
+            std::string::npos);
+}
+
+TEST(XlaCompilationTest,
+     DynamicExpressionCompatibilityRejectsDifferentVariableSets) {
+  const std::vector<xla::DExpr> exprs = {
+      xla::DExpr::Var(1) + xla::DExpr::Var(2),
+      xla::DExpr::Var(1) + xla::DExpr::Var(3)};
+
+  const std::optional<std::string> incompatibility =
+      testing::CheckDynamicExpressionCompatibilityForTest(exprs);
+
+  ASSERT_TRUE(incompatibility.has_value());
+  EXPECT_NE(incompatibility->find("same variable ids"), std::string::npos);
+}
+
+TEST(XlaCompilationTest,
+     DynamicExpressionCompatibilityUsesCompleteRepeatedVariableCore) {
+  const xla::DExpr a = xla::DExpr::Var(1);
+  const xla::DExpr b = xla::DExpr::Var(2);
+  const xla::DExpr shared_core = (a + b) * (a - b);
+  const std::vector<xla::DExpr> exprs = {shared_core + 2, shared_core + 3};
+
+  EXPECT_FALSE(
+      testing::CheckDynamicExpressionCompatibilityForTest(exprs).has_value());
+}
+
+TEST(XlaCompilationTest,
+     DynamicExpressionCompatibilityAcceptsSharedCoreWithDifferentScaling) {
+  const xla::DExpr shared_core =
+      xla::DExpr::Var(1) + xla::DExpr::Var(2);
+  const std::vector<xla::DExpr> exprs = {
+      shared_core / 240, shared_core, xla::DExpr::Const(3) * shared_core};
+
+  EXPECT_FALSE(
+      testing::CheckDynamicExpressionCompatibilityForTest(exprs).has_value());
+}
+
+TEST(XlaCompilationTest,
+     DynamicExpressionCompatibilityAcceptsSingleVariableExpressions) {
+  const std::vector<xla::DExpr> exprs = {
+      xla::DExpr::Var(1) + 2,
+      xla::DExpr::Var(1) * xla::DExpr::Const(3)};
+
+  EXPECT_FALSE(
+      testing::CheckDynamicExpressionCompatibilityForTest(exprs).has_value());
+}
+
 TEST(XlaCompilationTest, Chains) {
   std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
   {
diff --git a/tensorflow/compiler/jit/partially_decluster_pass.cc b/tensorflow/compiler/jit/partially_decluster_pass.cc
index 442635c9a29696..bdc123a9e3d9d3 100644
--- a/tensorflow/compiler/jit/partially_decluster_pass.cc
+++ b/tensorflow/compiler/jit/partially_decluster_pass.cc
@@ -19,8 +19,10 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/strings/str_cat.h"
 #include "tensorflow/compiler/jit/device_util.h"
+#include "tensorflow/compiler/jit/encapsulate_util.h"
 #include "tensorflow/compiler/jit/xla_cluster_util.h"
 #include "tensorflow/compiler/tf2xla/const_analysis.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "tensorflow/core/common_runtime/function.h"
 #include "tensorflow/core/framework/function.h"
@@ -47,6 +49,10 @@ absl::Status FindNodesToDecluster(const Graph& graph,
   MemoryTypeVector input_mtypes, output_mtypes;
 
   for (Node* n : post_order) {
+    if (SymbolicContentEnabled() &&
+        n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) {
+      continue;
+    }
     std::optional<absl::string_view> from_cluster = GetXlaClusterForNode(*n);
     if (!from_cluster) {
       continue;
@@ -308,6 +314,10 @@ absl::Status PartiallyDeclusterGraph(Graph* graph,
     if (!compile_time_const_nodes[n->id()]) {
       continue;
     }
+    if (SymbolicContentEnabled() &&
+        n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) {
+      continue;
+    }
 
     absl::string_view cluster_name = *GetXlaClusterForNode(*n);
     bool node_on_cluster_edge =
@@ -379,6 +389,10 @@ absl::Status PartiallyDeclusterGraph(Graph* graph) {
     if (!IsShapeConsumerOp(*n)) {
       continue;
     }
+    if (SymbolicContentEnabled() &&
+        n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) {
+      continue;
+    }
 
     std::optional<absl::string_view> cluster = GetXlaClusterForNode(*n);
     if (!cluster.has_value()) {
diff --git a/tensorflow/compiler/jit/shape_inference.cc b/tensorflow/compiler/jit/shape_inference.cc
index 7a5106aa69bbbf..b447b6802a0607 100644
--- a/tensorflow/compiler/jit/shape_inference.cc
+++ b/tensorflow/compiler/jit/shape_inference.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "absl/log/log.h"
 #include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
 #include "tensorflow/core/common_runtime/shape_refiner.h"
 #include "tensorflow/core/framework/function.h"
@@ -49,10 +50,29 @@ absl::Status ShapeHandleToTensorShape(
   if (!context->RankKnown(handle)) return absl::OkStatus();
 
   std::vector<int64_t> dims(context->Rank(handle));
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
+  std::vector<xla::DExpr> dyn_exprs;
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    dyn_exprs.resize(context->Rank(handle));
+  }
   for (int32_t i = 0, end = dims.size(); i < end; ++i) {
     dims[i] = context->Value(context->Dim(handle, i));
+    if (flags->tf_xla_enable_dynamic_sizes) {
+      DimExpr* expr = context->GetDimExpr(context->Dim(handle, i));
+      dyn_exprs[i] = expr != nullptr
+                         ? *expr
+                         : context->ValueKnown(context->Dim(handle, i))
+                               ? xla::DExpr::Const(dims[i])
+                               : xla::DExpr::Unknown(
+                                     xla::kMissingExpressionSentinel);
+    }
+  }
+  auto status =
+      PartialTensorShape::MakePartialShape(dims.data(), dims.size(), shape);
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    shape->set_expressions(std::move(dyn_exprs));
   }
-  return PartialTensorShape::MakePartialShape(dims.data(), dims.size(), shape);
+  return status;
 }
 
 absl::Status PropagateShapes(
diff --git a/tensorflow/compiler/jit/xla_batch_matcher.cc b/tensorflow/compiler/jit/xla_batch_matcher.cc
new file mode 100644
index 00000000000000..ba7b539500aa2b
--- /dev/null
+++ b/tensorflow/compiler/jit/xla_batch_matcher.cc
@@ -0,0 +1,268 @@
+/* Copyright 2026 The TensorFlow Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/jit/xla_batch_matcher.h"
+
+#include <algorithm>
+#include <cctype>
+#include <cstdlib>
+#include <sstream>
+#include <stdexcept>
+
+#include "absl/strings/numbers.h"
+#include "xla/debug_options_flags.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+
+namespace {
+
+// Trims spaces and tabs from both ends of a string.
+std::string Trim(const std::string& value) {
+  size_t start = value.find_first_not_of(" \t");
+  size_t end = value.find_last_not_of(" \t");
+  return start == std::string::npos ? ""
+                                    : value.substr(start, end - start + 1);
+}
+
+}  // namespace
+
+XlaBatchMatcher::XlaBatchMatcher() {
+  env_str_ = xla::GetDebugOptionsFromFlags().xla_compile_batch_sizes();
+  parse_env_config();
+}
+
+std::vector<int64_t> XlaBatchMatcher::parse_single_item(const std::string& item) {
+  std::vector<int64_t> batch_list;
+  if (item.empty()) return batch_list;
+
+  // 1. Parse single value (no colon separator)
+  if (item.find(':') == std::string::npos) {
+    // Validate all characters are digits (reject non-numeric chars)
+    for (char c : item) {
+      if (!isdigit(c)) {
+        throw std::invalid_argument("Non-numeric characters: " + item);
+      }
+    }
+    auto val = static_cast<int64_t>(std::stoll(item));
+    if (val <= 0 || val > kMaxBatch) {
+      throw std::invalid_argument("Out of valid range (1-" + std::to_string(kMaxBatch) + "): " + item);
+    }
+    batch_list.push_back(val);
+    return batch_list;
+  }
+
+  // 2. Parse range format (start:end:step)
+  std::stringstream ss(item);
+  std::string part;
+  std::vector<std::string> parts;
+  while (std::getline(ss, part, ':')) {
+    parts.push_back(Trim(part));
+  }
+  if (parts.size() > 3) {
+    throw std::invalid_argument("Invalid range format (requires start:end:step): " + item);
+  }
+
+  // Convert and validate range parameters
+  int64_t start, end, step;
+  try {
+    start = std::stoi(parts[0]);
+    end = std::stoi(parts[1]);
+    step = parts.size() == 2 ? 1 : std::stoi(parts[2]);
+  } catch (...) {
+    throw std::invalid_argument("Invalid numeric values in range: " + item);
+  }
+  if (start <= 0 || end <= 0 || step <= 0) {
+    throw std::invalid_argument("Range parameters must be positive integers: " + item);
+  }
+  if (start > end) {
+    throw std::invalid_argument("Start value > end value in range: " + item);
+  }
+  if (end > kMaxBatch) {
+    throw std::invalid_argument("Range exceeds max limit (" + std::to_string(kMaxBatch) + "): " + item);
+  }
+
+  // Generate batch list from range
+  for (int64_t i = start; i <= end; i += step) {
+    batch_list.push_back(i);
+  }
+  return batch_list;
+}
+
+void XlaBatchMatcher::print_all_batches(const std::string& cluster_key,
+                                       const std::vector<int64_t>& batches) {
+  std::ostringstream oss;
+  oss << "[XLA_BATCH_INFO] cluster_key=" << cluster_key
+      << " valid batch list update: ";
+  for (size_t i = 0; i < batches.size(); ++i) {
+    if (i > 0) oss << ", ";
+    oss << batches[i];
+  }
+  LOG(INFO) << oss.str();
+}
+
+// Parse environment variable config into a seed batch list used to initialize
+// each cluster-key state.
+void XlaBatchMatcher::parse_env_config() {
+  // If the env var not set or is empty, leave seed empty and let per-cluster
+  // logic generate factor-preserving padded batches.
+  if (env_str_.empty()) {
+    VLOG(2) << "[XLA_BATCH_WARN] Env var " << "--tf_xla_compile_batch_sizes" <<
+        "is empty, will use factor-preserving padding by default";
+    return;
+  }
+
+  // Parse into the default seed state.
+  ClusterState& seed = clusters_[""];
+
+  std::stringstream ss(env_str_);
+  std::string item;
+  while (std::getline(ss, item, ',')) {
+    std::string trimmed_item = Trim(item);
+    if (trimmed_item.empty()) continue;
+
+    try {
+      std::vector<int64_t> item_batches = parse_single_item(trimmed_item);
+      seed.all_batches.insert(seed.all_batches.end(), item_batches.begin(),
+                              item_batches.end());
+    } catch (const std::exception& e) {
+      LOG(INFO) << "[XLA_BATCH_WARN] Failed to parse config item, skipping: "
+                << trimmed_item << " (" << e.what() << ")";
+    }
+  }
+
+  if (!seed.all_batches.empty()) {
+    std::sort(seed.all_batches.begin(), seed.all_batches.end());
+    auto last = std::unique(seed.all_batches.begin(), seed.all_batches.end());
+    seed.all_batches.erase(last, seed.all_batches.end());
+    print_all_batches("<seed>", seed.all_batches);
+  }
+}
+
+// Calculate the smallest power of two greater than the real batch
+static int64_t GetNextPowerOfTwo(int64_t real_batch) {
+  // If real_batch is already a power of two, return real_batch directly
+  if ((real_batch & (real_batch - 1)) == 0) {
+    return real_batch;
+  }
+
+  int64_t power = 1;
+  while (power < real_batch) {
+    power <<= 1;
+  }
+  return power;
+}
+
+// Small/large multipliers for padding are configurable via environment variables:
+// TF_XLA_BATCH_SMALL_FACTOR (default 10) and TF_XLA_BATCH_LARGE_FACTOR (default 2).
+static int64_t GetEnvFactorOrDefault(const char* name, int64_t def) {
+  const char* val = std::getenv(name);
+  if (val == nullptr) return def;
+  int64_t parsed = 0;
+  if (!absl::SimpleAtoi(val, &parsed) || parsed <= 0) {
+    LOG(WARNING) << "[XLA_BATCH_WARN] Failed to parse env var " << name
+                 << "=\"" << val << "\", using default: " << def;
+    return def;
+  }
+  return parsed;
+}
+
+static int64_t GetFactorPreservingBatch(int64_t real_batch) {
+  // Cache env values in function-local statics (thread-safe since C++11).
+  static const int64_t small_factor =
+      GetEnvFactorOrDefault("TF_XLA_BATCH_SMALL_FACTOR", 10);
+  static const int64_t large_factor =
+      GetEnvFactorOrDefault("TF_XLA_BATCH_LARGE_FACTOR", 2);
+
+  const int64_t k = (real_batch < 10) ? small_factor : large_factor;
+  // Guard overflow / max range.
+  if (real_batch > kMaxBatch / k) {
+    LOG(WARNING) << "[XLA_BATCH_ERR] Out of valid range: " << real_batch;
+    return real_batch;
+  }
+  return real_batch * k;
+}
+
+int64_t XlaBatchMatcher::find_min_larger_batch(ClusterState* state,
+                                              int64_t real_batch) {
+  if (real_batch <= 0 || real_batch > kMaxBatch) {
+    LOG(INFO) << "[XLA_BATCH_WARN] Out of valid range: " << real_batch;
+    return real_batch;
+  }
+
+  // 1) If there is no configured candidate list, do not use power-of-two.
+  if (state->all_batches.empty()) {
+    const int val = GetFactorPreservingBatch(real_batch);
+    state->all_batches.push_back(val);
+    state->updated = true;
+    return val;
+  }
+
+  // 2) Prefer a configured candidate strictly larger than real_batch.
+  auto ub = std::upper_bound(state->all_batches.begin(), state->all_batches.end(),
+                             real_batch);
+  if (ub != state->all_batches.end()) {
+    return *ub;
+  }
+
+  // 3) real_batch > all_batches_.back(): generate a factor-preserving padded
+  // batch, write it back, keep list sorted/unique.
+  const int64_t val = GetFactorPreservingBatch(real_batch);
+  auto insert_pos =
+      std::lower_bound(state->all_batches.begin(), state->all_batches.end(), val);
+  if (insert_pos == state->all_batches.end() || *insert_pos != val) {
+    state->all_batches.insert(insert_pos, val);
+    state->updated = true;
+  }
+  return val;
+}
+
+int64_t XlaBatchMatcher::get_xla_compile_batch(const std::string& cluster_key,
+                                              int64_t real_batch) {
+  // Lazily initialize per-key state from the seed ("" key) to preserve previous
+  // behavior when an env list is provided.
+  ClusterState& state = clusters_[cluster_key];
+  state.updated = false;
+  if (state.all_batches.empty()) {
+    auto it = clusters_.find("");
+    if (it != clusters_.end()) {
+      state.all_batches = it->second.all_batches;
+    }
+  }
+
+  const int64_t selected = find_min_larger_batch(&state, real_batch);
+
+  if (state.updated) {
+    VLOG(2) << "[XLA_BATCH_INFO] cluster_key=" << cluster_key
+            << " real batch: " << real_batch
+            << " -> selected compile batch: " << selected;
+    print_all_batches(cluster_key, state.all_batches);
+  }
+
+  return selected;
+}
+
+int64_t XlaBatchMatcher::get_xla_compile_batch(int64_t real_batch) {
+  return get_xla_compile_batch("", real_batch);
+}
+
+std::vector<int64_t> XlaBatchMatcher::get_all_batches(
+    const std::string& cluster_key) {
+  auto it = clusters_.find(cluster_key);
+  if (it == clusters_.end()) return {};
+  return it->second.all_batches;
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/xla_batch_matcher.h b/tensorflow/compiler/jit/xla_batch_matcher.h
new file mode 100644
index 00000000000000..08476116aaf336
--- /dev/null
+++ b/tensorflow/compiler/jit/xla_batch_matcher.h
@@ -0,0 +1,65 @@
+/* Copyright 2026 The TensorFlow Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_JIT_XLA_BATCH_MATCHER_H_
+#define TENSORFLOW_COMPILER_JIT_XLA_BATCH_MATCHER_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+
+namespace tensorflow {
+
+// Define the maximum allowed batch size: 2147483648 >> 1 = 1073741824
+// Exceeding this value will make the next power of two exceed the safe range
+constexpr int kMaxBatch = 2147483648ULL >> 1;
+
+class XlaBatchMatcher {
+ public:
+  XlaBatchMatcher();
+  virtual ~XlaBatchMatcher() = default;
+
+  // Per-cluster-key API (recommended): isolates padding candidates per cluster.
+  int64_t get_xla_compile_batch(const std::string& cluster_key,
+                                int64_t real_batch);
+
+  // Backward-compatible API: uses a default (global) key.
+  int64_t get_xla_compile_batch(int64_t real_batch);
+
+  // For debugging.
+  std::vector<int64_t> get_all_batches(const std::string& cluster_key);
+
+ private:
+  struct ClusterState {
+    std::vector<int64_t> all_batches;
+    bool updated = false;
+  };
+
+  void parse_env_config();
+  void print_all_batches(const std::string& cluster_key,
+                         const std::vector<int64_t>& batches);
+  std::vector<int64_t> parse_single_item(const std::string& item);
+
+  int64_t find_min_larger_batch(ClusterState* state, int64_t real_batch);
+
+  absl::flat_hash_map<std::string, ClusterState> clusters_;
+  std::string env_str_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_JIT_XLA_BATCH_MATCHER_H_
diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc
index 588161df309131..b6b1f08fe73dcb 100644
--- a/tensorflow/compiler/jit/xla_launch_util.cc
+++ b/tensorflow/compiler/jit/xla_launch_util.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <memory>
 #include <optional>
 #include <set>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -26,6 +27,7 @@ limitations under the License.
 #include "absl/cleanup/cleanup.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #include "tensorflow/compiler/jit/pjrt_tensor_buffer.h"
 #include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h"
@@ -38,6 +40,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_common.h"
 #include "xla/pjrt/pjrt_future.h"
+#include "xla/printer.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
 #include "xla/stream_executor/platform_manager.h"
@@ -47,6 +50,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/gpu/gpu_serving_device_selector.h"
 #include "tensorflow/core/common_runtime/gpu_device_context.h"
 #include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/batch_size_resource.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/resource_mgr.h"
@@ -63,6 +67,17 @@ limitations under the License.
 #include "tsl/platform/statusor.h"
 
 namespace tensorflow {
+
+std::string DExprToString(const xla::DExpr& expr) {
+  xla::DExpr simplified = expr.simplify();
+  if (!simplified && !simplified.is_unknown()) {
+    return "";
+  }
+  xla::StringPrinter printer;
+  simplified->print(&printer);
+  return std::move(printer).ToString();
+}
+
 namespace {
 using xla::ScopedShapedBuffer;
 using xla::ShapedBuffer;
@@ -361,7 +376,8 @@ absl::Status XlaComputationLaunchContext::PopulateOutputs(
     ScopedShapedBuffer output, int missing_ctx_input_prefix,
     absl::Span<VariableInfo> variable_infos,
     const xla::HloInputOutputAliasConfig& input_output_alias,
-    const std::map<int, const Tensor*>& resource_vars) {
+    const std::map<int, const Tensor*>& resource_vars,
+    const xla::ExecutableRunOptions* run_options) {
   se::Stream* stream =
       ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr;
   Allocator* allocator = ctx->device()->GetAllocator({});
@@ -430,10 +446,123 @@ absl::Status XlaComputationLaunchContext::PopulateOutputs(
     }
   } else {
     for (int i = 0; i < ctx->num_outputs(); ++i) {
-      output_tensor_shapes.push_back(compilation_result->outputs[i].shape);
+      xla::Shape output_host_shape = output.on_host_shape();
+      const xla::Shape& subshape =
+          xla::ShapeUtil::GetSubshape(output_host_shape, {i});
+      VLOG(2) << "PopulateOutputs: subshape[" << i << "]: "<< subshape;
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(XLAShapeToTensorShape(subshape, &shape));
+      bool has_dynamic = false;
+
+      for (int dim = 0; dim < subshape.expressions().size(); ++dim) {
+        const auto& expr = subshape.expressions(dim);
+        if (expr && expr->is_dynamic()) {
+          has_dynamic = true;
+          VLOG(1) << "Current expression is " << expr;
+          if (run_options) {
+            const int64_t run_options_batch_size = run_options->batch_size();
+            if (run_options_batch_size <= 0) {
+              return absl::InvalidArgumentError(absl::StrCat(
+                  "Cannot substitute dynamic output shape for output ", i,
+                  ", dimension ", dim,
+                  ": the XLA runtime batch size was not initialized"));
+            }
+            VLOG(1) << "PopulateOutputs read run_options->batch_size()="
+                      << run_options_batch_size << " for output " << i
+                      << " dimension " << dim;
+            xla::DExpr batch_size = xla::DExpr::Const(run_options_batch_size);
+            const std::set<int> ids = expr->get_all_ids();
+            if (ids.size() != 1) {
+              return absl::InvalidArgumentError(absl::StrCat(
+                  "Runtime shape substitution expected exactly one dynamic "
+                  "variable for output ",
+                  i, ", dimension ", dim, ", but found ", ids.size(),
+                  " variables in expression ", DExprToString(expr)));
+            }
+            const int substitute_var_id = *ids.begin();
+            VLOG(1) << "Calling output shape substitute with run_options for "
+                      << "output " << i << " dimension " << dim
+                      << " expr=" << DExprToString(expr)
+                      << " substitute Var(" << substitute_var_id << ")="
+                      << DExprToString(batch_size);
+            xla::DExpr subst_expr =
+                expr.substitute(substitute_var_id, batch_size).simplify();
+            VLOG(1) << "Output shape substitute with run_options for output "
+                      << i << " dimension " << dim << " returned "
+                      << DExprToString(subst_expr);
+            if (!subst_expr->is_constant()) {
+              return absl::InvalidArgumentError(absl::StrCat(
+                  "Runtime shape substitution did not produce an integer "
+                  "constant for output ",
+                  i, ", dimension ", dim, ": ", DExprToString(subst_expr)));
+            }
+            shape.set_dim(dim, subst_expr->get_val());
+          } else {
+            // TODO: Fallback to BatchSizeResource for now. Remove it later.
+            LOG(WARNING) << "PopulateOutputs did not receive run_options for "
+                         << "output " << i << " dimension " << dim
+                         << "; falling back to BatchSizeResource";
+            BatchSizeResource* bsr = nullptr;
+            ScopedStepContainer* step_container = ctx->step_container();
+            TF_RETURN_IF_ERROR(step_container->Lookup<BatchSizeResource>(
+                          ctx->resource_manager(), BatchSizeResourceName, &bsr));
+            if (bsr == nullptr) {
+              return errors::Internal(
+                  "BatchSizeResource lookup succeeded but returned null");
+            }
+            core::ScopedUnref bsr_ref(bsr);
+            const int64_t runtime_batch_size = bsr->GetBatchSize();
+            if (runtime_batch_size <= 0) {
+              return absl::InvalidArgumentError(absl::StrCat(
+                  "Cannot substitute dynamic output shape for output ", i,
+                  ", dimension ", dim,
+                  ": the XLA runtime batch size was not initialized"));
+            }
+            xla::DExpr batch_size = xla::DExpr::Const(runtime_batch_size);
+            const std::set<int> ids = expr->get_all_ids();
+            if (ids.size() != 1) {
+              return absl::InvalidArgumentError(absl::StrCat(
+                  "Runtime shape substitution expected exactly one dynamic "
+                  "variable for output ",
+                  i, ", dimension ", dim, ", but found ", ids.size(),
+                  " variables in expression ", DExprToString(expr)));
+            }
+            const int substitute_var_id = *ids.begin();
+            VLOG(1) << "Calling output shape substitute with "
+                      << "BatchSizeResource for output " << i
+                      << " dimension " << dim
+                      << " expr=" << DExprToString(expr)
+                      << " substitute Var(" << substitute_var_id
+                      << ")=" << DExprToString(batch_size);
+            xla::DExpr subst_expr =
+                expr.substitute(substitute_var_id, batch_size).simplify();
+            VLOG(1) << "Output shape substitute with BatchSizeResource for "
+                      << "output " << i << " dimension " << dim
+                      << " returned " << DExprToString(subst_expr);
+            if (!subst_expr->is_constant()) {
+              return absl::InvalidArgumentError(absl::StrCat(
+                  "Runtime shape substitution did not produce an integer "
+                  "constant for output ",
+                  i, ", dimension ", dim, ": ", DExprToString(subst_expr)));
+            }
+            shape.set_dim(dim, subst_expr->get_val());
+          }
+        }
+      }
+      if (has_dynamic) {
+        output_tensor_shapes.push_back(shape);
+      }
+      else {
+        output_tensor_shapes.push_back(compilation_result->outputs[i].shape);
+      }
     }
   }
 
+  VLOG(2) << "output_tensor_shapes:";
+  for (auto s:output_tensor_shapes) {
+    VLOG(2) << s;
+  }
+
   // Copy XLA results to the OpOutputList.
   int output_num = 0;
   for (int i = 0, end = ctx->num_outputs(); i < end; ++i) {
diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h
index 5e5128d515bf97..e7b4c462b306cd 100644
--- a/tensorflow/compiler/jit/xla_launch_util.h
+++ b/tensorflow/compiler/jit/xla_launch_util.h
@@ -18,6 +18,7 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_
 #define TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_
 
+#include <string>
 #include <map>
 #include <memory>
 #include <set>
@@ -29,14 +30,21 @@ limitations under the License.
 #include "xla/client/local_client.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/service/shaped_buffer.h"
+#include "xla/executable_run_options.h"
 #include "xla/stream_executor/device_memory_allocator.h"
 #include "tensorflow/core/framework/allocation_description.pb.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 
+namespace xla {
+class DExpr;
+}
+
 namespace tensorflow {
 
+std::string DExprToString(const xla::DExpr& expr);
+
 // Creates a list of updated resource variables.
 absl::StatusOr<std::vector<VariableInfo>> GatherVariableInfo(
     OpKernelContext* ctx,
@@ -208,7 +216,8 @@ class XlaComputationLaunchContext {
       xla::ScopedShapedBuffer output, int missing_ctx_input_prefix,
       absl::Span<VariableInfo> variable_infos,
       const xla::HloInputOutputAliasConfig& input_output_alias,
-      const std::map<int, const Tensor*>& resource_vars);
+      const std::map<int, const Tensor*>& resource_vars,
+      const xla::ExecutableRunOptions* run_options = nullptr);
 
  private:
   xla::LocalClient* client_;
diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index f2080a0752f41c..8432ed574b7b62 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -739,6 +739,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/types:optional",
+        "@local_tsl//tsl/platform:protobuf",
         "@local_xla//xla/client",
         "@local_xla//xla/hlo/builder:value_inference",
         "@local_xla//xla/hlo/builder:xla_builder",
@@ -851,6 +852,7 @@ cc_library(
         ":common",
         ":xla_argument",
         ":xla_helpers",
+        "//tensorflow/compiler/jit:flags",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
@@ -874,10 +876,12 @@ cc_library(
     hdrs = [
         "literal_util.h",
         "shape_util.h",
+        "symbolic_content_util.h",
         "type_util.h",
     ],
     visibility = [":friends"],
     deps = [
+        "//tensorflow/compiler/jit:flags",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
@@ -1092,6 +1096,8 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
         "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/grappler:grappler_item",
+        "//tensorflow/core/grappler/costs:graph_properties",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest",
@@ -1563,6 +1569,7 @@ cc_library(
     deps = [
         ":xla_compiler",
         ":xla_expression",
+        "//tensorflow/compiler/jit:flags",
         "//tensorflow/compiler/jit:xla_compile_util",
         "//tensorflow/compiler/mlir/tf2xla/api/v1:compile_mlir_util_no_tf_dialect_passes",
         "//tensorflow/compiler/mlir/utils:array_container_utils",
diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD
index cd11a497b83514..37d68dbd23c242 100644
--- a/tensorflow/compiler/tf2xla/kernels/BUILD
+++ b/tensorflow/compiler/tf2xla/kernels/BUILD
@@ -1504,6 +1504,7 @@ tf_kernel_library(
     name = "slice_op",
     srcs = ["slice_op.cc"],
     deps = [
+        "//tensorflow/compiler/tf2xla:common",
         "//tensorflow/compiler/tf2xla:xla_compilation_device",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_context",
@@ -2200,6 +2201,7 @@ tf_kernel_library(
         "//tensorflow/compiler/tf2xla/ops:xla_ops",
         "//tensorflow/core:framework",
         "//tensorflow/core:protos_all_cc",
+        "@local_tsl//tsl/platform:protobuf",
         "@local_xla//xla/hlo/builder:xla_builder",
     ],
 )
@@ -2210,6 +2212,7 @@ tf_kernel_library(
     deps = [
         ":shape_util",
         ":tensor_list_utils",
+        "//tensorflow/compiler/tf2xla:common",
         "//tensorflow/compiler/tf2xla:xla_compilation_device",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_context",
diff --git a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
index 0dd528e3dea173..1db3ab4c78c461 100644
--- a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc
@@ -21,9 +21,9 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/compiler/tf2xla/kernels/relu_op.h"
-#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/constants.h"
@@ -241,7 +241,9 @@ class FusedBatchNormOpEx : public FusedBatchNormOp {
 
 REGISTER_XLA_OP(Name("FusedBatchNorm"), FusedBatchNormOp);
 REGISTER_XLA_OP(Name("FusedBatchNormV2"), FusedBatchNormOp);
-REGISTER_XLA_OP(Name("FusedBatchNormV3"), MlirXlaOpKernel);
+REGISTER_XLA_OP_FACTORY(
+    Name("FusedBatchNormV3"),
+    CreateDynamicNativeXlaOpKernel<FusedBatchNormOpV3>);
 REGISTER_XLA_OP(Name("_FusedBatchNormEx"), FusedBatchNormOpEx);
 
 class FusedBatchNormGradOp : public XlaOpKernel {
@@ -358,7 +360,9 @@ class FusedBatchNormGradOp : public XlaOpKernel {
 
 REGISTER_XLA_OP(Name("FusedBatchNormGrad"), FusedBatchNormGradOp);
 REGISTER_XLA_OP(Name("FusedBatchNormGradV2"), FusedBatchNormGradOp);
-REGISTER_XLA_OP(Name("FusedBatchNormGradV3"), MlirXlaOpKernel);
+REGISTER_XLA_OP_FACTORY(
+    Name("FusedBatchNormGradV3"),
+    CreateDynamicNativeXlaOpKernel<FusedBatchNormGradOp>);
 
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc
index 7a42150f3a9c19..6deeaa7ba11037 100644
--- a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc
@@ -36,6 +36,8 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
   const int input_rank = input_tensor_shape.dims();
   const absl::InlinedVector<int64_t, 4> input_shape =
       input_tensor_shape.dim_sizes();
+  const std::vector<xla::DExpr> input_exprs =
+      input_tensor_shape.get_filled_expressions();
   const int block_rank = block_shape.size();
 
   OP_REQUIRES(
@@ -76,11 +78,21 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
                               ") is not divisible by product of block sizes (",
                               block_num_elems, ")"));
   std::vector<int64_t> reshaped_shape(input_rank + block_rank);
+  std::vector<xla::DExpr> reshaped_exprs(input_rank + block_rank);
   std::copy(block_shape.begin(), block_shape.end(), reshaped_shape.begin());
+  std::fill(reshaped_exprs.begin(), reshaped_exprs.begin() + block_rank,
+            xla::DExpr::Const(0));
+  for (int i = 0; i < block_rank; ++i) {
+    reshaped_exprs[i] = xla::DExpr::Const(block_shape[i]);
+  }
   reshaped_shape[block_rank] = batch_size / block_num_elems;
+  reshaped_exprs[block_rank] =
+      (input_exprs[0] / xla::DExpr::Const(block_num_elems)).simplify();
   std::copy(input_shape.begin() + 1, input_shape.end(),
             reshaped_shape.begin() + block_rank + 1);
-  xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape);
+  std::copy(input_exprs.begin() + 1, input_exprs.end(),
+            reshaped_exprs.begin() + block_rank + 1);
+  xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs);
 
   // 2. Permute dimensions of `reshaped` to produce `permuted` of shape
   //      [batch / prod(block_shape),
@@ -111,15 +123,22 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
   //       ...,
   //       input_shape[N-1]]
   std::vector<int64_t> reshaped_permuted_shape(input_rank);
+  std::vector<xla::DExpr> reshaped_permuted_exprs(input_rank);
   reshaped_permuted_shape[0] = batch_size / block_num_elems;
+  reshaped_permuted_exprs[0] =
+      (input_exprs[0] / xla::DExpr::Const(block_num_elems)).simplify();
   for (int i = 0; i < block_rank; ++i) {
     reshaped_permuted_shape[1 + i] = block_shape[i] * input_shape[1 + i];
+    reshaped_permuted_exprs[1 + i] =
+        (xla::DExpr::Const(block_shape[i]) * input_exprs[1 + i]).simplify();
   }
   std::copy(remainder_shape.begin(), remainder_shape.end(),
             reshaped_permuted_shape.begin() + 1 + block_rank);
+  std::copy(input_exprs.begin() + 1 + block_rank, input_exprs.end(),
+            reshaped_permuted_exprs.begin() + 1 + block_rank);
 
   xla::XlaOp reshaped_permuted =
-      xla::Reshape(permuted, reshaped_permuted_shape);
+      xla::Reshape(permuted, reshaped_permuted_shape, reshaped_permuted_exprs);
 
   // 4. Crop the start and end of dimensions `[1, ..., M]` of
   //    `reshaped_permuted` according to `crops` to produce the output of shape:
@@ -133,6 +152,9 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
   std::vector<int64_t> start_indices(input_rank, 0);
   std::vector<int64_t> end_indices = reshaped_permuted_shape;
   std::vector<int64_t> strides(input_rank, 1);
+  std::vector<xla::DExpr> start_exprs(input_rank, xla::DExpr::Const(0));
+  std::vector<xla::DExpr> end_exprs(reshaped_permuted_exprs.begin(),
+                                    reshaped_permuted_exprs.end());
   for (int i = 0; i < block_rank; ++i) {
     int64_t crop_start = crops.Get<int64_t>({i, 0});
     int64_t crop_end = crops.Get<int64_t>({i, 1});
@@ -140,14 +162,16 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp input,
                 errors::InvalidArgument("Crops must be non-negative"));
     start_indices[1 + i] = crop_start;
     end_indices[1 + i] -= crop_end;
+    start_exprs[1 + i] = xla::DExpr::Const(crop_start);
+    end_exprs[1 + i] = (reshaped_permuted_exprs[1 + i] - crop_end).simplify();
     OP_REQUIRES(
         ctx, start_indices[1 + i] <= end_indices[1 + i],
         errors::InvalidArgument(
             "Cropped size must be non-negative: start: ", crop_start,
             " end: ", crop_end, " size ", reshaped_permuted_shape[1 + i]));
   }
-  xla::XlaOp output =
-      xla::Slice(reshaped_permuted, start_indices, end_indices, strides);
+  xla::XlaOp output = xla::Slice(reshaped_permuted, start_indices, end_indices,
+                                 start_exprs, end_exprs, strides);
   ctx->SetOutput(0, output);
 }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/beta_op.cc b/tensorflow/compiler/tf2xla/kernels/beta_op.cc
index 4ead9f76fcee11..40cdd190d173b4 100644
--- a/tensorflow/compiler/tf2xla/kernels/beta_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/beta_op.cc
@@ -63,11 +63,14 @@ class BetaincOp : public XlaOpKernel {
     auto result =
         builder->ReportErrorOrReturn([&]() -> absl::StatusOr<xla::XlaOp> {
           TF_ASSIGN_OR_RETURN(
-              auto a, BroadcastTo(ctx->Input(0), merged_shape.dim_sizes()));
+              auto a, BroadcastTo(ctx->Input(0), merged_shape.dim_sizes(),
+                                  merged_shape.get_filled_expressions()));
           TF_ASSIGN_OR_RETURN(
-              auto b, BroadcastTo(ctx->Input(1), merged_shape.dim_sizes()));
+              auto b, BroadcastTo(ctx->Input(1), merged_shape.dim_sizes(),
+                                  merged_shape.get_filled_expressions()));
           TF_ASSIGN_OR_RETURN(
-              auto x, BroadcastTo(ctx->Input(2), merged_shape.dim_sizes()));
+              auto x, BroadcastTo(ctx->Input(2), merged_shape.dim_sizes(),
+                                  merged_shape.get_filled_expressions()));
           return xla::RegularizedIncompleteBeta(a, b, x);
         });
     ctx->SetOutput(0, result);
diff --git a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
index e9f571d830d619..34762dfeb9b545 100644
--- a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc
@@ -42,34 +42,46 @@ namespace {
 // A subclass of a XlaBinaryOp must build the computation that
 // describes the (tensor,tensor)->tensor function to apply to each element of
 // the input.
-#define XLA_MAKE_BINARY(NAME, HLO)                                         \
+#define XLA_MAKE_BINARY(NAME, HLO, SYMBOLIC_HLO)                           \
   class NAME##Op : public XlaBinaryOp {                                    \
    public:                                                                 \
     explicit NAME##Op(OpKernelConstruction* ctx) : XlaBinaryOp(ctx) {}     \
     xla::XlaOp Computation(                                                \
         XlaOpKernelContext* ctx, const xla::XlaOp& lhs,                    \
-        const absl::Span<const int64_t>& lhs_shape, const xla::XlaOp& rhs, \
+        const absl::Span<const int64_t>& lhs_shape,                        \
+        const xla::XlaOp& rhs,                                             \
         const absl::Span<const int64_t>& rhs_shape,                        \
         const BCast& broadcast_helper,                                     \
+        const absl::Span<const xla::DExpr>& broadcast_output_exprs,        \
         const std::vector<int64_t>& extend_dimensions) override {          \
       xla::XlaBuilder* b = ctx->builder();                                 \
       (void)b;                                                             \
       (void)lhs_shape;                                                     \
       (void)rhs_shape;                                                     \
+      (void)broadcast_output_exprs;                                        \
       (void)extend_dimensions;                                             \
       return HLO;                                                          \
     }                                                                      \
+    xla::DExpr SymbolicComputation(const xla::DExpr& lhs,                  \
+                                   const xla::DExpr& rhs) override {       \
+      return SYMBOLIC_HLO;                                                 \
+    }                                                                      \
   };                                                                       \
   REGISTER_XLA_OP(Name(#NAME), NAME##Op)
 
-XLA_MAKE_BINARY(Add, xla::Add(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(AddV2, xla::Add(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Sub, xla::Sub(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Mul, xla::Mul(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Div, xla::Div(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Add, xla::Add(lhs, rhs, extend_dimensions),
+                (lhs + rhs).simplify());
+XLA_MAKE_BINARY(AddV2, xla::Add(lhs, rhs, extend_dimensions),
+                (lhs + rhs).simplify());
+XLA_MAKE_BINARY(Sub, xla::Sub(lhs, rhs, extend_dimensions),
+                (lhs - rhs).simplify());
+XLA_MAKE_BINARY(Mul, xla::Mul(lhs, rhs, extend_dimensions),
+                (lhs * rhs).simplify());
+XLA_MAKE_BINARY(Div, xla::Div(lhs, rhs, extend_dimensions),
+                (lhs / rhs).simplify());
 
-XLA_MAKE_BINARY(Atan2, xla::Atan2(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Atan2, xla::Atan2(lhs, rhs, extend_dimensions), xla::DExpr());
+XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions), xla::DExpr());
 
 // Implementation of DivNoNan. Pseudo-code:
 // if (y == 0) {
@@ -78,8 +90,12 @@ XLA_MAKE_BINARY(Complex, xla::Complex(lhs, rhs, extend_dimensions));
 //   return x / y;
 // }
 static xla::XlaOp DivNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
-                               xla::XlaOp y, const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+                               xla::XlaOp y,
+                               const absl::Span<const xla::DExpr>&
+                                   broadcast_output_exprs,
+                               const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   auto zero = XlaHelpers::Zero(b, dtype);
   auto y_equals_0 = xla::Eq(y, zero);
   auto zeros = xla::ZerosLike(x);
@@ -87,7 +103,9 @@ static xla::XlaOp DivNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
   return result;
 }
 XLA_MAKE_BINARY(DivNoNan,
-                DivNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+                DivNoNanImpl(b, input_type(0), lhs, rhs,
+                             broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 // Implementation of MulNoNan. Pseudo-code:
 // if (y == 0) {
@@ -96,8 +114,12 @@ XLA_MAKE_BINARY(DivNoNan,
 //   return x * y;
 // }
 static xla::XlaOp MulNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
-                               xla::XlaOp y, const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+                               xla::XlaOp y,
+                               const absl::Span<const xla::DExpr>&
+                                   broadcast_output_exprs,
+                               const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   auto zero = XlaHelpers::Zero(b, dtype);
   auto y_equals_0 = xla::Eq(y, zero);
   auto zeros = xla::ZerosLike(x);
@@ -105,7 +127,9 @@ static xla::XlaOp MulNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
   return result;
 }
 XLA_MAKE_BINARY(MulNoNan,
-                MulNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+                MulNoNanImpl(b, input_type(0), lhs, rhs,
+                             broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 // Implementation of FloorDiv.
 //
@@ -118,8 +142,12 @@ XLA_MAKE_BINARY(MulNoNan,
 //   return z;
 // }
 static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
-                               xla::XlaOp y, const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+                               xla::XlaOp y,
+                               const absl::Span<const xla::DExpr>&
+                                   broadcast_output_exprs,
+                               const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   if (DataTypeIsFloating(dtype)) {
     if (dtype == DataType::DT_BFLOAT16) {
       // The result of a BF16 division may produce the Ceil of what was
@@ -144,43 +172,61 @@ static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
   return xla::Select(round_down, xla::Sub(x_div_y, one), x_div_y);
 }
 XLA_MAKE_BINARY(FloorDiv,
-                FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));
+                FloorDivImpl(b, input_type(0), lhs, rhs,
+                             broadcast_output_exprs, broadcast_helper),
+                (lhs / rhs).simplify());
 
 xla::XlaOp XlogyImpl(xla::XlaOp x, xla::XlaOp y,
+                     const absl::Span<const xla::DExpr>& broadcast_output_exprs,
                      const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   auto zero = xla::ZerosLike(x);
   auto is_zero = xla::Eq(x, zero);
   return xla::Select(is_zero, zero, xla::Mul(x, xla::Log(y)));
 }
-XLA_MAKE_BINARY(Xlogy, XlogyImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(Xlogy,
+                XlogyImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 xla::XlaOp Xlog1pyImpl(xla::XlaOp x, xla::XlaOp y,
+                       const absl::Span<const xla::DExpr>& broadcast_output_exprs,
                        const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   auto non_zero = xla::Mul(x, xla::Log1p(y));
   auto zero = xla::ZerosLike(non_zero);
   auto x_is_zero = xla::Eq(x, zero);
   return xla::Select(x_is_zero, zero, non_zero);
 }
-XLA_MAKE_BINARY(Xlog1py, Xlog1pyImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(Xlog1py,
+                Xlog1pyImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 xla::XlaOp XdivyImpl(xla::XlaOp x, xla::XlaOp y,
+                     const absl::Span<const xla::DExpr>& broadcast_output_exprs,
                      const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   auto zero = xla::ZerosLike(x);
   auto is_zero = xla::Eq(x, zero);
   return xla::Select(is_zero, zero, xla::Div(x, y));
 }
-XLA_MAKE_BINARY(Xdivy, XdivyImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(Xdivy,
+                XdivyImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 // Implementation of FloorMod. Pseudo-code:
 // T trunc_mod = std::fmod(x, y);
 // return trunc_mod != 0 && (y < 0 != trunc_mod < 0) ? trunc_mod + y
 //                                                   : trunc_mod;
 static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
-                               xla::XlaOp y, const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+                               xla::XlaOp y,
+                               const absl::Span<const xla::DExpr>&
+                                   broadcast_output_exprs,
+                               const BCast& broadcast_helper) {
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   auto zero = XlaHelpers::Zero(b, dtype);
   auto trunc_mod = xla::Rem(x, y);
   auto trunc_mod_not_zero = xla::Ne(trunc_mod, zero);
@@ -189,43 +235,62 @@ static xla::XlaOp FloorModImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
   return xla::Select(do_plus, xla::Add(trunc_mod, y), trunc_mod);
 }
 XLA_MAKE_BINARY(FloorMod,
-                FloorModImpl(b, input_type(0), lhs, rhs, broadcast_helper));
-
-XLA_MAKE_BINARY(BitwiseAnd, xla::And(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(BitwiseOr, xla::Or(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(BitwiseXor, xla::Xor(lhs, rhs, extend_dimensions));
-
-XLA_MAKE_BINARY(LeftShift, xla::ShiftLeft(lhs, rhs, extend_dimensions));
+                FloorModImpl(b, input_type(0), lhs, rhs,
+                             broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
+
+XLA_MAKE_BINARY(BitwiseAnd, xla::And(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(BitwiseOr, xla::Or(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(BitwiseXor, xla::Xor(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+
+XLA_MAKE_BINARY(LeftShift, xla::ShiftLeft(lhs, rhs, extend_dimensions),
+                xla::DExpr());
 XLA_MAKE_BINARY(RightShift,
                 (DataTypeIsUnsigned(ctx->input_type(0))
                      ? xla::ShiftRightLogical(lhs, rhs, extend_dimensions)
-                     : xla::ShiftRightArithmetic(lhs, rhs, extend_dimensions)));
-
-XLA_MAKE_BINARY(LogicalAnd, xla::And(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(LogicalOr, xla::Or(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Mod, xla::Rem(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Maximum, xla::Max(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Minimum, xla::Min(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(RealDiv, xla::Div(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(ReciprocalGrad, xla::Neg(xla::Mul(rhs, xla::Mul(lhs, lhs))));
+                     : xla::ShiftRightArithmetic(lhs, rhs, extend_dimensions)),
+                xla::DExpr());
+
+XLA_MAKE_BINARY(LogicalAnd, xla::And(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(LogicalOr, xla::Or(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(Mod, xla::Rem(lhs, rhs, extend_dimensions), xla::DExpr());
+XLA_MAKE_BINARY(Maximum, xla::Max(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(Minimum, xla::Min(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(RealDiv, xla::Div(lhs, rhs, extend_dimensions),
+                (lhs / rhs).simplify());
+XLA_MAKE_BINARY(ReciprocalGrad, xla::Neg(xla::Mul(rhs, xla::Mul(lhs, lhs))),
+                xla::DExpr());
 XLA_MAKE_BINARY(
     RsqrtGrad,
     xla::Mul((lhs * lhs) * lhs,
              xla::Div(rhs, XlaHelpers::IntegerLiteral(b, input_type(0), -2)),
-             extend_dimensions));
+             extend_dimensions),
+    xla::DExpr());
 XLA_MAKE_BINARY(
     SqrtGrad,
     xla::Div(xla::Mul(rhs, XlaHelpers::FloatLiteral(b, input_type(0), 0.5)),
-             lhs, extend_dimensions));
+             lhs, extend_dimensions),
+    xla::DExpr());
 
 // Implementation of TruncateDiv.
 //
 // For floating-point values, returns trunc(x / y).  For integers, simply
 // returns x / y.
 static xla::XlaOp TruncateDivImpl(xla::XlaBuilder* b, DataType dtype,
-                                  xla::XlaOp x, xla::XlaOp y,
+                                  xla::XlaOp x,
+                                  xla::XlaOp y,
+                                  const absl::Span<const xla::DExpr>&
+                                      broadcast_output_exprs,
                                   const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   if (!DataTypeIsFloating(dtype)) {
     return xla::Div(x, y);
   }
@@ -235,35 +300,46 @@ static xla::XlaOp TruncateDivImpl(xla::XlaBuilder* b, DataType dtype,
   return xla::Select(round_up, xla::Ceil(x_div_y), xla::Floor(x_div_y));
 }
 XLA_MAKE_BINARY(TruncateDiv,
-                TruncateDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));
-XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions));
+                TruncateDivImpl(b, input_type(0), lhs, rhs,
+                                broadcast_output_exprs, broadcast_helper),
+                (lhs / rhs).simplify());
+XLA_MAKE_BINARY(TruncateMod, xla::Rem(lhs, rhs, extend_dimensions),
+                xla::DExpr());
 
 // Comparison ops
-XLA_MAKE_BINARY(Equal, xla::Eq(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(NotEqual, xla::Ne(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Greater, xla::Gt(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(GreaterEqual, xla::Ge(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(Less, xla::Lt(lhs, rhs, extend_dimensions));
-XLA_MAKE_BINARY(LessEqual, xla::Le(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Equal, xla::Eq(lhs, rhs, extend_dimensions), xla::DExpr());
+XLA_MAKE_BINARY(NotEqual, xla::Ne(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(Greater, xla::Gt(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(GreaterEqual, xla::Ge(lhs, rhs, extend_dimensions),
+                xla::DExpr());
+XLA_MAKE_BINARY(Less, xla::Lt(lhs, rhs, extend_dimensions), xla::DExpr());
+XLA_MAKE_BINARY(LessEqual, xla::Le(lhs, rhs, extend_dimensions),
+                xla::DExpr());
 
 // Non-linear ops
 XLA_MAKE_BINARY(SigmoidGrad,
                 xla::Mul(xla::Mul(rhs, lhs),
-                         xla::Sub(XlaHelpers::One(b, input_type(0)), lhs)));
+                         xla::Sub(XlaHelpers::One(b, input_type(0)), lhs)),
+                xla::DExpr());
 
-XLA_MAKE_BINARY(SoftplusGrad, xla::Mul(lhs, xla::Logistic(rhs)));
+XLA_MAKE_BINARY(SoftplusGrad, xla::Mul(lhs, xla::Logistic(rhs)),
+                xla::DExpr());
 
 // softsigngrad(gradients, features) = gradients / (1 + abs(features)) ** 2
 XLA_MAKE_BINARY(SoftsignGrad,
                 xla::Div(lhs,
                          xla::Square(xla::Add(XlaHelpers::One(b, input_type(0)),
-                                              xla::Abs(rhs)))));
+                                              xla::Abs(rhs)))),
+                xla::DExpr());
 
 XLA_MAKE_BINARY(TanhGrad,
                 xla::Mul(rhs, xla::Sub(XlaHelpers::One(b, input_type(0)),
-                                       xla::Mul(lhs, lhs))));
+                                       xla::Mul(lhs, lhs))),
+                xla::DExpr());
 
-XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions));
+XLA_MAKE_BINARY(Pow, xla::Pow(lhs, rhs, extend_dimensions), xla::DExpr());
 
 xla::XlaOp SquaredDifferenceImpl(
     DataType dtype, xla::XlaOp x, xla::XlaOp y,
@@ -277,55 +353,86 @@ xla::XlaOp SquaredDifferenceImpl(
 }
 XLA_MAKE_BINARY(SquaredDifference,
                 SquaredDifferenceImpl(input_type(0), lhs, rhs,
-                                      extend_dimensions));
+                                      extend_dimensions),
+                xla::DExpr());
 
 xla::XlaOp IgammaImpl(xla::XlaOp x, xla::XlaOp y,
+                      const absl::Span<const xla::DExpr>& broadcast_output_exprs,
                       const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   return xla::Igamma(x, y);
 }
 
-XLA_MAKE_BINARY(Igamma, IgammaImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(Igamma,
+                IgammaImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 xla::XlaOp IgammaGradAImpl(xla::XlaOp x, xla::XlaOp y,
+                           const absl::Span<const xla::DExpr>&
+                               broadcast_output_exprs,
                            const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   return xla::IgammaGradA(x, y);
 }
 
-XLA_MAKE_BINARY(IgammaGradA, IgammaGradAImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(IgammaGradA,
+                IgammaGradAImpl(lhs, rhs, broadcast_output_exprs,
+                                broadcast_helper),
+                xla::DExpr());
 
 xla::XlaOp RandomGammaGradImpl(xla::XlaOp x, xla::XlaOp y,
+                               const absl::Span<const xla::DExpr>&
+                                   broadcast_output_exprs,
                                const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   return xla::RandomGammaGrad(x, y);
 }
 
 XLA_MAKE_BINARY(RandomGammaGrad,
-                RandomGammaGradImpl(lhs, rhs, broadcast_helper));
+                RandomGammaGradImpl(lhs, rhs, broadcast_output_exprs,
+                                    broadcast_helper),
+                xla::DExpr());
 
 xla::XlaOp IgammacImpl(xla::XlaOp x, xla::XlaOp y,
+                       const absl::Span<const xla::DExpr>& broadcast_output_exprs,
                        const BCast& broadcast_helper) {
-  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
+  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper,
+                                          broadcast_output_exprs);
   return xla::Igammac(x, y);
 }
 
-XLA_MAKE_BINARY(Igammac, IgammacImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(Igammac,
+                IgammacImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 xla::XlaOp PolygammaImpl(xla::XlaOp n, xla::XlaOp x,
+                         const absl::Span<const xla::DExpr>&
+                             broadcast_output_exprs,
                          const BCast& broadcast_helper) {
-  std::tie(n, x) = XlaBinaryOp::Broadcast(n, x, broadcast_helper);
+  std::tie(n, x) = XlaBinaryOp::Broadcast(n, x, broadcast_helper,
+                                          broadcast_output_exprs);
   return xla::Polygamma(n, x);
 }
 
-XLA_MAKE_BINARY(Polygamma, PolygammaImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(Polygamma,
+                PolygammaImpl(lhs, rhs, broadcast_output_exprs,
+                              broadcast_helper),
+                xla::DExpr());
 
-xla::XlaOp ZetaImpl(xla::XlaOp x, xla::XlaOp q, const BCast& broadcast_helper) {
-  std::tie(x, q) = XlaBinaryOp::Broadcast(x, q, broadcast_helper);
+xla::XlaOp ZetaImpl(xla::XlaOp x, xla::XlaOp q,
+                    const absl::Span<const xla::DExpr>& broadcast_output_exprs,
+                    const BCast& broadcast_helper) {
+  std::tie(x, q) = XlaBinaryOp::Broadcast(x, q, broadcast_helper,
+                                          broadcast_output_exprs);
   return xla::Zeta(x, q);
 }
 
-XLA_MAKE_BINARY(Zeta, ZetaImpl(lhs, rhs, broadcast_helper));
+XLA_MAKE_BINARY(Zeta,
+                ZetaImpl(lhs, rhs, broadcast_output_exprs, broadcast_helper),
+                xla::DExpr());
 
 #undef XLA_MAKE_BINARY
 
diff --git a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
index fadd2f87219464..5e277772f40e4e 100644
--- a/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/bincount_op.cc
@@ -110,16 +110,24 @@ class DenseBincountOp : public XlaOpKernel {
     scatter_dnums.add_scatter_dims_to_operand_dims(0);
 
     if (rank == 2) {
-      output_shape = xla::ShapeUtil::MakeShape(dtype, {size, output_size});
+      output_shape = xla::ShapeUtil::MakeShape(
+          dtype, {size, output_size},
+          std::vector<xla::DExpr>{input_shape.expressions(0),
+                                  xla::DExpr::Const(output_size)});
       scatter_dnums.add_inserted_window_dims(1);
       scatter_dnums.add_scatter_dims_to_operand_dims(1);
-      auto i_shape =
-          xla::ShapeUtil::MakeShape(input_xla_type, {input_shape.dimensions()});
+      auto i_shape = xla::ShapeUtil::MakeShape(input_xla_type,
+                                               input_shape.dimensions(),
+                                               input_shape.expressions());
       auto i = xla::Iota(ctx->builder(), i_shape, 0);
+      xla::DExpr flattened_expr =
+          input_shape.expressions(0) * input_shape.expressions(1);
       i = xla::Reshape(
-          i, {input_shape.dimensions(0) * input_shape.dimensions(1), 1});
+          i, {input_shape.dimensions(0) * input_shape.dimensions(1), 1},
+          {flattened_expr, xla::DExpr::Const(1)});
       auto j = xla::Reshape(
-          input, {input_shape.dimensions(0) * input_shape.dimensions(1), 1});
+          input, {input_shape.dimensions(0) * input_shape.dimensions(1), 1},
+          {flattened_expr, xla::DExpr::Const(1)});
       std::vector<xla::XlaOp> iotas_to_concat;
       iotas_to_concat.push_back(i);
       iotas_to_concat.push_back(j);
@@ -127,10 +135,12 @@ class DenseBincountOp : public XlaOpKernel {
       updates = xla::Broadcast(
           one, {input_shape.dimensions(0) * input_shape.dimensions(1)});
       output = xla::Broadcast(
-          zero, {output_shape.dimensions(0), output_shape.dimensions(1)});
+          zero, {output_shape.dimensions(0), output_shape.dimensions(1)},
+          {output_shape.expressions(0), output_shape.expressions(1)});
       if (has_weights && !binary_output_) {
         weights = xla::Reshape(
-            weights, {input_shape.dimensions(0) * input_shape.dimensions(1)});
+            weights, {input_shape.dimensions(0) * input_shape.dimensions(1)},
+            {flattened_expr});
         updates = weights;
       }
     } else {
diff --git a/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc
index 975179466bf104..d3e76ee63da1c3 100644
--- a/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc
@@ -37,7 +37,8 @@ class BroadcastToOp : public XlaOpKernel {
                    context->ConstantInputAsShape(
                        1, &output_shape, xla::ValueInferenceMode::kUpperBound));
     auto output_status_or =
-        BroadcastTo(context->Input(0), output_shape.dim_sizes());
+        BroadcastTo(context->Input(0), output_shape.dim_sizes(),
+                    output_shape.get_filled_expressions());
     OP_REQUIRES_OK(context, output_status_or.status());
     auto output = output_status_or.value();
     std::vector<bool> dynamic_dims;
diff --git a/tensorflow/compiler/tf2xla/kernels/cast_op.cc b/tensorflow/compiler/tf2xla/kernels/cast_op.cc
index 1779cfcc1ced40..74e9c5aa4003ff 100644
--- a/tensorflow/compiler/tf2xla/kernels/cast_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/cast_op.cc
@@ -92,6 +92,15 @@ class CastOp : public XlaOpKernel {
       output = xla::ConvertElementType(input, dst_type_);
     }
 
+    const auto& input_contents = ctx->InputExpression(0).contents();
+    if (!input_contents.empty()) {
+      auto output_expr =
+          XlaExpression::XlaOp(output, ctx->expected_output_dtype(0));
+      output_expr.set_contents(std::vector<xla::DExpr>(input_contents.begin(),
+                                                       input_contents.end()));
+      ctx->SetOutputExpression(0, output_expr);
+      return;
+    }
     ctx->SetOutput(0, output);
   }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc
index 6b4f278c72beff..928fbb6e22d76f 100644
--- a/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/clip_by_value_op.cc
@@ -46,11 +46,11 @@ class ClipByValueOp : public XlaOpKernel {
 
     if (shape != min_shape) {
       OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(min_shape), shape_error());
-      min = xla::Broadcast(min, shape.dim_sizes());
+      min = xla::Broadcast(min, shape.dim_sizes(), shape.get_filled_expressions());
     }
     if (shape != max_shape) {
       OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(max_shape), shape_error());
-      max = xla::Broadcast(max, shape.dim_sizes());
+      max = xla::Broadcast(max, shape.dim_sizes(), shape.get_filled_expressions());
     }
     ctx->SetOutput(0, xla::Clamp(min, input, max));
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/concat_op.cc b/tensorflow/compiler/tf2xla/kernels/concat_op.cc
index bed3479941ca41..8a71a7dc4d1128 100644
--- a/tensorflow/compiler/tf2xla/kernels/concat_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/concat_op.cc
@@ -18,9 +18,11 @@ limitations under the License.
 #include <cstdint>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/log/log.h"
 #include "absl/strings/str_join.h"
 #include "tensorflow/compiler/tf2xla/kernels/shape_util.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
@@ -41,6 +43,44 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+bool AppendConcatInputContents(const XlaExpression& expr,
+                               const TensorShape& shape,
+                               std::vector<xla::DExpr>* contents) {
+  const auto& input_contents = expr.contents();
+  if (!input_contents.empty()) {
+    if (shape.dims() == 0) {
+      if (input_contents.size() != 1) {
+        return false;
+      }
+      contents->push_back(input_contents[0]);
+      return true;
+    }
+    if (shape.dims() == 1) {
+      contents->insert(contents->end(), input_contents.begin(),
+                       input_contents.end());
+      return true;
+    }
+    return false;
+  }
+  if (shape.dims() == 0) {
+    contents->push_back(xla::DExpr::Unknown(xla::kUnknownContentSentinel));
+    return true;
+  }
+  if (shape.dims() != 1) {
+    return false;
+  }
+  for (int64_t i = 0; i < shape.dim_size(0); ++i) {
+    contents->push_back(xla::DExpr::Unknown(xla::kUnknownContentSentinel));
+  }
+  return true;
+}
+
+bool HasDynamicContents(absl::Span<const xla::DExpr> contents) {
+  return absl::c_any_of(contents, [](const xla::DExpr& expr) {
+    return expr && expr->is_dynamic();
+  });
+}
+
 // --------------------------------------------------------------------------
 class ConcatBaseOp : public XlaOpKernel {
  public:
@@ -74,6 +114,29 @@ class ConcatBaseOp : public XlaOpKernel {
     // Make a vector holding the XlaOp for each of the inputs that has non-zero
     // elements.
     std::vector<xla::XlaOp> input_data;
+    std::vector<std::vector<xla::DExpr>> input_contents;
+    input_contents.resize(N);
+    bool has_output_contents = SymbolicContentEnabled() && axis == 0;
+    std::vector<xla::DExpr> output_contents;
+    if (has_output_contents) {
+      for (int i = 0; i < N; ++i) {
+        if (!AppendConcatInputContents(ctx->InputExpression(ValueInputIndex(i)),
+                                       shapes[i],
+                                       &output_contents)) {
+          has_output_contents = false;
+          output_contents.clear();
+          break;
+        }
+        if (shapes[i].dims() == 0) {
+          input_contents[i].push_back(output_contents.back());
+        } else if (shapes[i].dims() == 1) {
+          input_contents[i].insert(input_contents[i].end(),
+                                   output_contents.end() - shapes[i].dim_size(0),
+                                   output_contents.end());
+        }
+      }
+      has_output_contents = has_output_contents && HasDynamicContents(output_contents);
+    }
     int output_concat_dim = 0;
     for (int i = 0; i < N; ++i) {
       xla::XlaOp handle = values[i];
@@ -86,7 +149,12 @@ class ConcatBaseOp : public XlaOpKernel {
               "] = ", in_shape.DebugString()));
       if (in_shape.dims() == 0) {
         // Inputs that come in as scalars must be reshaped to 1-vectors.
-        input_data.push_back(xla::Reshape(handle, {1}));
+        xla::XlaOp reshaped = xla::Reshape(handle, {1});
+        if (has_output_contents) {
+          OP_REQUIRES_OK(
+              ctx, ctx->builder()->SetInstructionContents(reshaped, input_contents[i]));
+        }
+        input_data.push_back(reshaped);
       } else {
         input_data.push_back(handle);
       }
@@ -94,10 +162,24 @@ class ConcatBaseOp : public XlaOpKernel {
     }
 
     VLOG(1) << "Concat dim " << concat_dim << " equivalent to " << axis;
-    ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), input_data, axis));
+    auto output = xla::ConcatInDim(ctx->builder(), input_data, axis);
+    if (has_output_contents) {
+      OP_REQUIRES_OK(
+          ctx, ctx->builder()->SetInstructionContents(output, output_contents));
+      auto output_expr =
+          XlaExpression::XlaOp(output, ctx->expected_output_dtype(0));
+      output_expr.set_contents(std::move(output_contents));
+      ctx->SetOutputExpression(0, output_expr);
+      return;
+    }
+    ctx->SetOutput(0, output);
   }
 
  private:
+  int ValueInputIndex(int value_index) const {
+    return value_index < axis_index_ ? value_index : value_index + 1;
+  }
+
   int axis_index_;
 };
 
diff --git a/tensorflow/compiler/tf2xla/kernels/const_op.cc b/tensorflow/compiler/tf2xla/kernels/const_op.cc
index d2463a9974b1bb..33d8b2079897d0 100644
--- a/tensorflow/compiler/tf2xla/kernels/const_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/const_op.cc
@@ -15,19 +15,29 @@ limitations under the License.
 
 #include <cstdint>
 #include <type_traits>
+#include <vector>
 
+#include "xla/printer.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape_expr.h"
+#include "tsl/platform/protobuf.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/logging.h"
 
 namespace tensorflow {
 namespace {
 
+constexpr char kUserInferredValueContentsAttrName[] =
+    "_user_inferred_value_contents";
+
 template <typename DstT, typename SrcT>
 DstT CastTo(SrcT src) {
   return static_cast<DstT>(src);
@@ -100,6 +110,40 @@ xla::XlaOp GetScalarConst(const TensorProto& proto, xla::XlaBuilder* b) {
   return xla::XlaOp();
 }
 
+std::vector<xla::DExpr> BuildShapeContentsFromTensorShapeProto(
+    const TensorShapeProto& shape) {
+  std::vector<xla::DExpr> contents;
+  contents.reserve(shape.dim_size());
+  for (int i = 0; i < shape.dim_size(); ++i) {
+    xla::DExpr expr;
+    if (i < shape.expressions_size()) {
+      expr = DimExprFromProto(shape.expressions(i));
+    }
+    if (expr) {
+      xla::StringPrinter printer;
+      expr->print(&printer);
+      VLOG(1) << "BuildShapeContentsFromTensorShape dim=" << i
+              << " expr=" << std::move(printer).ToString()
+              << " dynamic=" << (expr->is_dynamic() ? "true" : "false");
+    } else {
+      VLOG(1) << "BuildShapeContentsFromTensorShape dim=" << i
+              << " expr=_ dynamic=false";
+    }
+    contents.push_back(expr && expr->is_dynamic()
+                           ? std::move(expr)
+                           : xla::DExpr::Unknown(
+                                 xla::kUnknownContentSentinel));
+  }
+  return contents;
+}
+
+bool CanAttachContentsFromTensorShapeProto(const TensorShape& tensor_shape,
+                                           const TensorShapeProto& contents) {
+  return (tensor_shape.dims() == 0 && contents.dim_size() == 1) ||
+         (tensor_shape.dims() == 1 &&
+          tensor_shape.dim_size(0) == contents.dim_size());
+}
+
 class ConstOp : public XlaOpKernel {
  public:
   explicit ConstOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
@@ -117,13 +161,49 @@ class ConstOp : public XlaOpKernel {
   void Compile(XlaOpKernelContext* ctx) override {
     xla::XlaBuilder* b = ctx->builder();
 
+    bool has_dynamic = false;
+    TensorShapeProto inferred_shape_proto;
+    TensorShapeProto inferred_value_contents_proto;
+    string inferred_value_contents_serialized;
+    if (GetNodeAttr(ctx->op_kernel().def(), "has_dynamic", &has_dynamic).ok() &&
+        has_dynamic) {
+      GetNodeAttr(ctx->op_kernel().def(), "user_inferred_shape",
+                  &inferred_shape_proto)
+          .IgnoreError();
+    }
+    if (GetNodeAttr(ctx->op_kernel().def(), kUserInferredValueContentsAttrName,
+                    &inferred_value_contents_serialized)
+            .ok()) {
+      if (!inferred_value_contents_proto.ParseFromString(
+              inferred_value_contents_serialized)) {
+        inferred_value_contents_proto.Clear();
+      }
+    }
+    const bool has_contents_proto = inferred_value_contents_proto.dim_size() > 0;
+    const TensorShapeProto& contents_proto =
+        has_contents_proto ? inferred_value_contents_proto : inferred_shape_proto;
+
     // To avoid blowups for large constants filled with the same value,
     // recognize that case and emit a scalar broadcast instead.
     TensorShape shape(proto_.tensor_shape());
     if (shape.num_elements() > 1) {
       xla::XlaOp value = GetScalarConst(proto_, b);
       if (value.valid()) {
-        ctx->SetOutput(0, xla::Broadcast(value, shape.dim_sizes()));
+        if (has_dynamic) {
+          VLOG(1) << "ConstOp broadcast fast path shape="
+                    << shape.DebugString() << " inferred_rank="
+                    << inferred_shape_proto.dim_size();
+        }
+        xla::XlaOp broadcast =
+            xla::Broadcast(value, shape.dim_sizes(), shape.get_expressions());
+        XlaExpression output =
+            XlaExpression::XlaOp(broadcast, ctx->expected_output_dtype(0));
+        if ((has_contents_proto || has_dynamic) &&
+            CanAttachContentsFromTensorShapeProto(shape, contents_proto)) {
+          output.set_contents(
+              BuildShapeContentsFromTensorShapeProto(contents_proto));
+        }
+        ctx->SetOutputExpression(0, output);
         return;
       }
     }
@@ -132,7 +212,13 @@ class ConstOp : public XlaOpKernel {
     OP_REQUIRES(ctx, tensor.FromProto(cpu_allocator(), proto_),
                 errors::InvalidArgument("Cannot parse tensor from proto: ",
                                         proto_.DebugString()));
-    ctx->SetConstantOutput(0, tensor);
+    XlaExpression output = XlaExpression::Constant(tensor);
+    if ((has_contents_proto || has_dynamic) &&
+        CanAttachContentsFromTensorShapeProto(tensor.shape(), contents_proto)) {
+      output.set_contents(
+          BuildShapeContentsFromTensorShapeProto(contents_proto));
+    }
+    ctx->SetOutputExpression(0, output);
   }
 
  private:
diff --git a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc
index 3fe22dcb4441e7..7889852a24c580 100644
--- a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc
+++ b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc
@@ -97,8 +97,11 @@ xla::XlaOp TransposeFilterForGroupConvolutionBackpropInput(
   CHECK_GE(num_dims, 2);  // Crash OK
   xla::Shape new_shape = filter_shape;
   new_shape.set_dimensions(num_dims - 1, num_groups);
-  new_shape.add_dimensions(filter_shape.dimensions(num_dims - 1) / num_groups);
-  xla::XlaOp result = xla::Reshape(filter, new_shape.dimensions());
+  new_shape.add_dimensions(
+      filter_shape.dimensions(num_dims - 1) / num_groups,
+      filter_shape.expressions(num_dims - 1) / num_groups);
+  xla::XlaOp result =
+      xla::Reshape(filter, new_shape.dimensions(), new_shape.expressions());
 
   // 2. Transpose to [H, W, ..., G, filter_in_depth, out_depth / G]
   std::vector<int64_t> transpose_dims(num_dims + 1);
@@ -118,7 +121,8 @@ xla::XlaOp ReshapeFilterForDepthwiseConvolution(const xla::Shape& filter_shape,
                                                 xla::XlaOp filter) {
   return xla::Reshape(
       filter,
-      GroupedFilterShapeForDepthwiseConvolution(filter_shape).dimensions());
+      GroupedFilterShapeForDepthwiseConvolution(filter_shape).dimensions(),
+      GroupedFilterShapeForDepthwiseConvolution(filter_shape).expressions());
 }
 
 // Performs some basic checks on ConvOpAttrs that are true for all kinds of XLA
@@ -603,7 +607,8 @@ absl::StatusOr<xla::XlaOp> MakeXlaBackpropFilterConvOp(
   }
 
   if (attrs.depthwise) {
-    filter_backprop = xla::Reshape(filter_backprop, filter_shape.dimensions());
+    filter_backprop = xla::Reshape(filter_backprop, filter_shape.dimensions(),
+                                   filter_shape.expressions());
   }
 
   return filter_backprop;
diff --git a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc
index b1da0acd61608f..f90bad207cfbd4 100644
--- a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc
@@ -110,7 +110,8 @@ class ConvNDOp : public XlaOpKernel {
         expanded_input_shape.set_dimensions(i + 1, input_shape.dimensions(i));
       }
       expanded_input_shape.set_dimensions(0, 1);
-      input = xla::Reshape(input, expanded_input_shape.dimensions());
+      input = xla::Reshape(input, expanded_input_shape.dimensions(),
+                           expanded_input_shape.expressions());
     } else if (attrs_.batch_dims > 1) {
       // Flatten batch_dims.
       std::vector<int64_t> to_collapse(attrs_.batch_dims);
@@ -131,7 +132,8 @@ class ConvNDOp : public XlaOpKernel {
     if (attrs_.batch_dims == 0) {
       xla::Shape no_batch_shape(out_shape);
       no_batch_shape.DeleteDimension(0);
-      out = xla::Reshape(out, no_batch_shape.dimensions());
+      out = xla::Reshape(out, no_batch_shape.dimensions(),
+                         no_batch_shape.expressions());
     } else if (attrs_.batch_dims > 1) {
       xla::Shape expanded_out_shape(input_shape);
       for (int i = attrs_.batch_dims; i < input_shape.dimensions().size();
@@ -139,7 +141,8 @@ class ConvNDOp : public XlaOpKernel {
         expanded_out_shape.set_dimensions(
             i, out_shape.dimensions(i - (attrs_.batch_dims - 1)));
       }
-      out = xla::Reshape(out, expanded_out_shape.dimensions());
+      out = xla::Reshape(out, expanded_out_shape.dimensions(),
+                         expanded_out_shape.expressions());
     }
     ctx->SetOutput(0, out);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc
index 6c91556862d9e2..a796509e36a4ba 100644
--- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc
@@ -32,10 +32,182 @@ limitations under the License.
 #include "xla/shape.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/bcast.h"
 
 namespace tensorflow {
 
+namespace {
+
+bool IsSymbolicContentType(DataType type) {
+  type = BaseType(type);
+  return type == DT_INT32 || type == DT_INT64;
+}
+
+bool TryGetIntContentsFromConstant(const Tensor& tensor,
+                                   std::vector<xla::DExpr>* contents) {
+  contents->clear();
+  if (tensor.dims() > 1) {
+    return false;
+  }
+  if (tensor.dtype() == DT_INT32) {
+    if (tensor.dims() == 0) {
+      contents->push_back(xla::DExpr::Const(tensor.scalar<int32>()()));
+      return true;
+    }
+    auto flat = tensor.flat<int32>();
+    contents->reserve(flat.size());
+    for (int i = 0; i < flat.size(); ++i) {
+      contents->push_back(xla::DExpr::Const(flat(i)));
+    }
+    return true;
+  }
+  if (tensor.dtype() == DT_INT64) {
+    if (tensor.dims() == 0) {
+      contents->push_back(xla::DExpr::Const(tensor.scalar<int64_t>()()));
+      return true;
+    }
+    auto flat = tensor.flat<int64_t>();
+    contents->reserve(flat.size());
+    for (int i = 0; i < flat.size(); ++i) {
+      contents->push_back(xla::DExpr::Const(flat(i)));
+    }
+    return true;
+  }
+  return false;
+}
+
+bool TryGetIntContentsFromLiteral(const xla::LiteralSlice& literal,
+                                  std::vector<xla::DExpr>* contents) {
+  contents->clear();
+  if (literal.shape().dimensions_size() > 1) {
+    return false;
+  }
+  if (literal.shape().element_type() == xla::S32) {
+    if (literal.shape().dimensions_size() == 0) {
+      contents->push_back(xla::DExpr::Const(literal.Get<int32>({})));
+      return true;
+    }
+    const int64_t size = literal.shape().dimensions(0);
+    contents->reserve(size);
+    for (int64_t i = 0; i < size; ++i) {
+      contents->push_back(xla::DExpr::Const(literal.Get<int32>({i})));
+    }
+    return true;
+  }
+  if (literal.shape().element_type() == xla::S64) {
+    if (literal.shape().dimensions_size() == 0) {
+      contents->push_back(xla::DExpr::Const(literal.Get<int64_t>({})));
+      return true;
+    }
+    const int64_t size = literal.shape().dimensions(0);
+    contents->reserve(size);
+    for (int64_t i = 0; i < size; ++i) {
+      contents->push_back(xla::DExpr::Const(literal.Get<int64_t>({i})));
+    }
+    return true;
+  }
+  return false;
+}
+
+bool TryGetInputContents(XlaOpKernelContext* ctx, const XlaExpression& expr,
+                         const TensorShape& shape,
+                         std::vector<xla::DExpr>* contents) {
+  if (shape.dims() > 1) {
+    return false;
+  }
+  if (!expr.contents().empty()) {
+    if (absl::c_any_of(expr.contents(),
+                       [](const xla::DExpr& e) { return !e; })) {
+      return false;
+    }
+    contents->assign(expr.contents().begin(), expr.contents().end());
+    return true;
+  }
+  auto constant = expr.constant_value();
+  if (!constant.has_value()) {
+    if (!expr.handle().valid() || expr.handle().IsUninitialized()) {
+      return false;
+    }
+    auto literal_or =
+        ctx->value_inference().AnalyzeConstant(expr.handle(),
+                                               xla::ValueInferenceMode::kValue);
+    if (!literal_or.ok() || !literal_or->AllValid()) {
+      return false;
+    }
+    auto literal = literal_or->GetValue();
+    if (!literal.has_value()) {
+      return false;
+    }
+    return TryGetIntContentsFromLiteral(*literal, contents);
+  }
+  return TryGetIntContentsFromConstant(*constant, contents);
+}
+
+xla::DExpr BroadcastedContentAt(absl::Span<const xla::DExpr> contents,
+                                const TensorShape& shape,
+                                int64_t output_index) {
+  if (shape.dims() == 0) {
+    return contents.empty() ? xla::DExpr() : contents[0];
+  }
+  if (contents.empty()) {
+    return xla::DExpr();
+  }
+  if (shape.dim_size(0) == 1) {
+    return contents[0];
+  }
+  if (output_index >= contents.size()) {
+    return xla::DExpr();
+  }
+  return contents[output_index];
+}
+
+bool TryBuildSymbolicBinaryContents(XlaOpKernelContext* ctx,
+                                    XlaBinaryOp* op,
+                                    const TensorShape& lhs_shape,
+                                    const TensorShape& rhs_shape,
+                                    const BCast& bcast,
+                                    std::vector<xla::DExpr>* contents) {
+  contents->clear();
+  if (!IsSymbolicContentType(ctx->input_type(0)) ||
+      !IsSymbolicContentType(ctx->expected_output_dtype(0))) {
+    return false;
+  }
+  const auto& output_shape = bcast.output_shape();
+  if (output_shape.size() > 1) {
+    return false;
+  }
+
+  std::vector<xla::DExpr> lhs_contents;
+  std::vector<xla::DExpr> rhs_contents;
+  if (!TryGetInputContents(ctx, ctx->InputExpression(0), lhs_shape,
+                           &lhs_contents) ||
+      !TryGetInputContents(ctx, ctx->InputExpression(1), rhs_shape,
+                           &rhs_contents)) {
+    return false;
+  }
+
+  int64_t output_elements = output_shape.empty() ? 1 : output_shape[0];
+  contents->reserve(output_elements);
+  for (int64_t i = 0; i < output_elements; ++i) {
+    xla::DExpr lhs_expr = BroadcastedContentAt(lhs_contents, lhs_shape, i);
+    xla::DExpr rhs_expr = BroadcastedContentAt(rhs_contents, rhs_shape, i);
+    if (!lhs_expr || !rhs_expr) {
+      contents->clear();
+      return false;
+    }
+    xla::DExpr out_expr = op->SymbolicComputation(lhs_expr, rhs_expr);
+    if (!out_expr) {
+      contents->clear();
+      return false;
+    }
+    contents->push_back(out_expr.simplify());
+  }
+  return true;
+}
+
+}  // namespace
+
 void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
   TensorShape lhs_shape = ctx->InputShape(0);
   TensorShape rhs_shape = ctx->InputShape(1);
@@ -68,6 +240,7 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
             lhs = xla::SliceInDim(lhs, 0, rhs_xla_shape.dimensions(i), 1,
                                   /*dimno=*/i);
             lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i));
+            lhs_tensor_shape->set_expression(i, rhs_xla_shape.expressions(i));
             // Propagate dynamic dimension.
             lhs = xla::SetDimensionSize(lhs, size, i);
           }
@@ -90,6 +263,7 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
                 lhs, xla::Zero(ctx->builder(), lhs_xla_shape.element_type()), i,
                 0, diff);
             lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i));
+            lhs_tensor_shape->set_expression(i, rhs_xla_shape.expressions(i));
             // Propagate dynamic dimension.
             lhs = xla::SetDimensionSize(lhs, size, i);
           }
@@ -137,6 +311,10 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
             lhs = xla::SetDimensionSize(lhs, size, i);
 
             lhs_tensor_shape->set_dim(i, rhs_xla_shape.dimensions(i));
+            lhs_tensor_shape->set_expression(
+                i, (lhs_tensor_shape->get_filled_expression(i) *
+                    rhs_xla_shape.expressions(i))
+                       .simplify());
           }
         }
       }
@@ -162,6 +340,67 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
     return;
   }
 
+  auto build_broadcast_output_expressions =
+      [&lhs_shape, &rhs_shape, &bcast]() -> std::vector<xla::DExpr> {
+    auto merge_broadcast_dim = [](bool has_lhs, int64_t lhs_dim,
+                                  const xla::DExpr& lhs_expr, bool has_rhs,
+                                  int64_t rhs_dim, const xla::DExpr& rhs_expr,
+                                  int64_t output_dim) {
+      if (!has_lhs) {
+        return rhs_expr;
+      }
+      if (!has_rhs) {
+        return lhs_expr;
+      }
+      if (lhs_dim == 1 && rhs_dim != 1) {
+        // A broadcasted singleton usually inherits the other side's
+        // expression, but keep a dynamic singleton visible by folding it into
+        // the result.
+        return lhs_expr && lhs_expr->is_dynamic()
+                   ? (lhs_expr * rhs_expr).simplify()
+                   : rhs_expr;
+      }
+      if (rhs_dim == 1 && lhs_dim != 1) {
+        // Symmetric case for a singleton rhs broadcast.
+        return rhs_expr && rhs_expr->is_dynamic()
+                   ? (rhs_expr * lhs_expr).simplify()
+                   : lhs_expr;
+      }
+      // When both sides describe the same logical dimension, prefer whichever
+      // side still carries a dynamic symbolic expression.
+      if (lhs_expr && lhs_expr->is_dynamic()) {
+        return lhs_expr;
+      }
+      if (rhs_expr && rhs_expr->is_dynamic()) {
+        return rhs_expr;
+      }
+      return xla::DExpr::Const(output_dim);
+    };
+
+    const auto& output_shape = bcast.output_shape();
+    std::vector<xla::DExpr> output_exprs(output_shape.size());
+
+    for (int out_i = output_shape.size() - 1, lhs_i = lhs_shape.dims() - 1,
+             rhs_i = rhs_shape.dims() - 1;
+         out_i >= 0; --out_i, --lhs_i, --rhs_i) {
+      const bool has_lhs = lhs_i >= 0;
+      const bool has_rhs = rhs_i >= 0;
+      xla::DExpr lhs_expr = has_lhs ? lhs_shape.get_filled_expression(lhs_i)
+                                    : xla::DExpr::Const(1);
+      xla::DExpr rhs_expr = has_rhs ? rhs_shape.get_filled_expression(rhs_i)
+                                    : xla::DExpr::Const(1);
+      const int64_t lhs_dim = has_lhs ? lhs_shape.dim_size(lhs_i) : 1;
+      const int64_t rhs_dim = has_rhs ? rhs_shape.dim_size(rhs_i) : 1;
+      output_exprs[out_i] = merge_broadcast_dim(
+          has_lhs, lhs_dim, lhs_expr, has_rhs, rhs_dim, rhs_expr,
+          output_shape[out_i]);
+    }
+
+    return output_exprs;
+  };
+  std::vector<xla::DExpr> broadcast_output_exprs =
+      build_broadcast_output_expressions();
+
   // If the ranks of the inputs don't match, TensorFlow automatically
   // reshapes the smaller by padding with dimensions of size 1 as a
   // prefix. In other words to pad a 5-vector to a 3-dimensional
@@ -187,22 +426,39 @@ void XlaBinaryOp::Compile(XlaOpKernelContext* ctx) {
   // Call virtual method to emit the computation.
   xla::XlaOp output =
       Computation(ctx, lhs_handle, lhs_shape.dim_sizes(), rhs_handle,
-                  rhs_shape.dim_sizes(), bcast, extend_dimension);
+                  rhs_shape.dim_sizes(), bcast,
+                  broadcast_output_exprs, extend_dimension);
 
   // The TensorFlow helper computed the post-broadcast shape in
   // output_shape: we rely on subclassed Computations to implement the
   // same broadcast semantics.
+  std::vector<xla::DExpr> output_contents;
+  if (TryBuildSymbolicBinaryContents(ctx, this, lhs_shape, rhs_shape, bcast,
+                                     &output_contents)) {
+    auto output_expr =
+        XlaExpression::XlaOp(output, ctx->expected_output_dtype(0));
+    output_expr.set_contents(std::move(output_contents));
+    ctx->SetOutputExpression(0, output_expr);
+    return;
+  }
   ctx->SetOutput(0, output);
 }
 
 /* static */ std::pair<xla::XlaOp, xla::XlaOp> XlaBinaryOp::Broadcast(
-    xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper) {
-  auto lhs_output = BroadcastTo(lhs, broadcast_helper.output_shape());
+    xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper,
+    absl::Span<const xla::DExpr> output_exprs) {
+  CHECK_EQ(output_exprs.size(), broadcast_helper.output_shape().size());
+  for (const xla::DExpr& expr : output_exprs) {
+    CHECK(expr);
+  }
+  auto lhs_output =
+      BroadcastTo(lhs, broadcast_helper.output_shape(), output_exprs);
   if (!lhs_output.ok()) {
     xla::XlaOp error = lhs.builder()->ReportError(lhs_output.status());
     return {error, error};
   }
-  auto rhs_output = BroadcastTo(rhs, broadcast_helper.output_shape());
+  auto rhs_output =
+      BroadcastTo(rhs, broadcast_helper.output_shape(), output_exprs);
   if (!rhs_output.ok()) {
     xla::XlaOp error = rhs.builder()->ReportError(rhs_output.status());
     return {error, error};
diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h
index d22e6eb74039b4..6e3734abe7d53f 100644
--- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h
+++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h
@@ -65,15 +65,25 @@ class XlaBinaryOp : public XlaOpKernel {
       XlaOpKernelContext* ctx, const xla::XlaOp& lhs,
       const absl::Span<const int64_t>& lhs_shape, const xla::XlaOp& rhs,
       const absl::Span<const int64_t>& rhs_shape, const BCast& broadcast_helper,
+      const absl::Span<const xla::DExpr>& broadcast_output_exprs,
       const std::vector<int64_t>& extend_dimensions) = 0;
 
+  // Returns a symbolic expression for one output element when content metadata
+  // should be propagated through this op. Returns an empty DExpr when the operation
+  // should not propagate symbolic contents.
+  virtual xla::DExpr SymbolicComputation(const xla::DExpr& lhs,
+                                         const xla::DExpr& rhs) {
+    return xla::DExpr();
+  }
+
   void Compile(XlaOpKernelContext* ctx) override;
 
   // Helper function that performs the broadcasting described by
   // 'broadcast_helper', yielding arguments 'lhs' and 'rhs' that have the same
   // shape.
   static std::pair<xla::XlaOp, xla::XlaOp> Broadcast(
-      xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper);
+      xla::XlaOp lhs, xla::XlaOp rhs, const BCast& broadcast_helper,
+      absl::Span<const xla::DExpr> output_exprs);
 };
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc
index e8e2babffd529c..3a651d1adae655 100644
--- a/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/depthtospace_op.cc
@@ -64,6 +64,7 @@ class DepthToSpaceOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx, input_xla_shape.status());
     absl::Span<const int64_t> input_shape =
         input_xla_shape.value().dimensions();
+    const xla::Shape& input_shape_with_exprs = input_xla_shape.value();
     int input_rank = input_shape.size();
 
     static const int kRequiredDims = 4;
@@ -77,20 +78,31 @@ class DepthToSpaceOp : public XlaOpKernel {
     std::vector<int64_t> reshaped_shape;
     std::vector<int64_t> transpose_order;
     std::vector<int64_t> output_shape;
+    std::vector<xla::DExpr> reshaped_exprs;
+    std::vector<xla::DExpr> output_exprs;
     reshaped_shape.reserve(input_rank);
     transpose_order.reserve(input_rank);
     output_shape.reserve(input_rank);
+    reshaped_exprs.reserve(input_rank + num_spatial_dims);
+    output_exprs.reserve(input_rank);
     if (data_format == FORMAT_NHWC) {
       reshaped_shape.push_back(input_shape[0]);
+      reshaped_exprs.push_back(input_shape_with_exprs.expressions(0));
       for (int i = 0; i < num_spatial_dims; ++i) {
         reshaped_shape.push_back(input_shape[1 + i]);
+        reshaped_exprs.push_back(input_shape_with_exprs.expressions(1 + i));
       }
       int64_t block_elems = 1;
       for (int i = 0; i < num_spatial_dims; ++i) {
         reshaped_shape.push_back(block_size_);
+        reshaped_exprs.push_back(xla::DExpr::Const(block_size_));
         block_elems *= block_size_;
       }
       reshaped_shape.push_back(input_shape[feature_dim] / block_elems);
+      reshaped_exprs.push_back(
+          (input_shape_with_exprs.expressions(feature_dim) /
+           xla::DExpr::Const(block_elems))
+              .simplify());
 
       transpose_order.push_back(0);
       for (int i = 0; i < num_spatial_dims; ++i) {
@@ -100,21 +112,37 @@ class DepthToSpaceOp : public XlaOpKernel {
       transpose_order.push_back(feature_dim + num_spatial_dims);
 
       output_shape.push_back(input_shape[0]);
+      output_exprs.push_back(input_shape_with_exprs.expressions(0));
       for (int i = 0; i < num_spatial_dims; ++i) {
         output_shape.push_back(input_shape[1 + i] * block_size_);
+        output_exprs.push_back(
+            (input_shape_with_exprs.expressions(1 + i) *
+             xla::DExpr::Const(block_size_))
+                .simplify());
       }
       output_shape.push_back(input_shape[feature_dim] / block_elems);
+      output_exprs.push_back(
+          (input_shape_with_exprs.expressions(feature_dim) /
+           xla::DExpr::Const(block_elems))
+              .simplify());
     } else {
       // NCHW format.
       reshaped_shape.push_back(input_shape[0]);
+      reshaped_exprs.push_back(input_shape_with_exprs.expressions(0));
       int64_t block_elems = 1;
       for (int i = 0; i < num_spatial_dims; ++i) {
         reshaped_shape.push_back(block_size_);
+        reshaped_exprs.push_back(xla::DExpr::Const(block_size_));
         block_elems *= block_size_;
       }
       reshaped_shape.push_back(input_shape[feature_dim] / block_elems);
+      reshaped_exprs.push_back(
+          (input_shape_with_exprs.expressions(feature_dim) /
+           xla::DExpr::Const(block_elems))
+              .simplify());
       for (int i = 0; i < num_spatial_dims; ++i) {
         reshaped_shape.push_back(input_shape[2 + i]);
+        reshaped_exprs.push_back(input_shape_with_exprs.expressions(2 + i));
       }
 
       transpose_order.push_back(0);
@@ -125,9 +153,18 @@ class DepthToSpaceOp : public XlaOpKernel {
       }
 
       output_shape.push_back(input_shape[0]);
+      output_exprs.push_back(input_shape_with_exprs.expressions(0));
       output_shape.push_back(input_shape[feature_dim] / block_elems);
+      output_exprs.push_back(
+          (input_shape_with_exprs.expressions(feature_dim) /
+           xla::DExpr::Const(block_elems))
+              .simplify());
       for (int i = 0; i < num_spatial_dims; ++i) {
         output_shape.push_back(input_shape[2 + i] * block_size_);
+        output_exprs.push_back(
+            (input_shape_with_exprs.expressions(2 + i) *
+             xla::DExpr::Const(block_size_))
+                .simplify());
       }
     }
 
@@ -148,7 +185,7 @@ class DepthToSpaceOp : public XlaOpKernel {
                     ") is not divisible by square of the block size (",
                     block_size_, ")"));
 
-    xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape);
+    xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs);
 
     // 2. Permute dimensions of `reshaped` to produce
     //    `permuted_reshaped` of shape:
@@ -169,7 +206,7 @@ class DepthToSpaceOp : public XlaOpKernel {
     //       input_shape[2] * block_size_,
     //       depth / (block_size_ * block_size_)]
     //
-    xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape);
+    xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape, output_exprs);
 
     // If this used to be a vectorized format turn it back now.
     if (data_format != data_format_) {
diff --git a/tensorflow/compiler/tf2xla/kernels/diag_op.cc b/tensorflow/compiler/tf2xla/kernels/diag_op.cc
index 404fa9f5e04e45..9644d52fa489a6 100644
--- a/tensorflow/compiler/tf2xla/kernels/diag_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/diag_op.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "xla/hlo/builder/lib/matrix.h"
 #include "xla/hlo/builder/lib/pooling.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape_util.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -38,7 +39,9 @@ namespace {
 
 // Create a diagonal / batch diagonal matrix with 'input' on the diagonal.
 xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size,
-                          absl::Span<const int64_t> other_dims) {
+                          const xla::DExpr& last_dim_expr,
+                          absl::Span<const int64_t> other_dims,
+                          absl::Span<const xla::DExpr> other_dim_exprs) {
   xla::XlaBuilder* builder = input.builder();
   // Create two matrices that have the following forms, and compare them:
   //
@@ -49,14 +52,23 @@ xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size,
   //
   // This produces a predicate matrix of the right size, with "true" on the
   // diagonal.
-  xla::XlaOp iota = xla::Iota(builder, xla::S32, last_dim_size);
-  xla::XlaOp iota_broadcast = xla::Broadcast(iota, {last_dim_size});
+  xla::XlaOp iota = xla::Iota(
+      builder,
+      xla::ShapeUtil::MakeShape(xla::S32, std::vector<int64_t>{last_dim_size},
+                                std::vector<xla::DExpr>{last_dim_expr}),
+      /*iota_dimension=*/0);
+  xla::XlaOp iota_broadcast = xla::Broadcast(
+      iota, {last_dim_size}, {last_dim_expr, last_dim_expr});
   xla::XlaOp mask = xla::Eq(iota_broadcast, iota, {0});
 
   // If this is a batched diagonal, broadcast the mask across the other
   // dimensions.
   if (!other_dims.empty()) {
-    mask = xla::Broadcast(mask, other_dims);
+    std::vector<xla::DExpr> mask_exprs(other_dim_exprs.begin(),
+                                       other_dim_exprs.end());
+    mask_exprs.push_back(last_dim_expr);
+    mask_exprs.push_back(last_dim_expr);
+    mask = xla::Broadcast(mask, other_dims, mask_exprs);
   }
 
   // Broadcast the input, and then use the mask computed above to select the
@@ -69,13 +81,17 @@ xla::XlaOp CreateDiagonal(xla::XlaOp input, int64_t last_dim_size,
   std::vector<int64_t> out_dim_sizes(other_dims.begin(), other_dims.end());
   out_dim_sizes.push_back(last_dim_size);
   out_dim_sizes.push_back(last_dim_size);
+  std::vector<xla::DExpr> out_dim_exprs(other_dim_exprs.begin(),
+                                        other_dim_exprs.end());
+  out_dim_exprs.push_back(last_dim_expr);
+  out_dim_exprs.push_back(last_dim_expr);
 
   // Broadcast into the second to last dimension.
   std::vector<int64_t> broadcast_dimensions(other_dims.size() + 1);
   absl::c_iota(broadcast_dimensions, 0);
   ++broadcast_dimensions.back();
-  xla::XlaOp input_broadcast =
-      xla::BroadcastInDim(input, out_dim_sizes, broadcast_dimensions);
+  xla::XlaOp input_broadcast = xla::BroadcastInDim(
+      input, out_dim_sizes, broadcast_dimensions, out_dim_exprs);
   return xla::Select(mask, input_broadcast, xla::ZerosLike(input_broadcast));
 }
 
@@ -102,17 +118,26 @@ class DiagOp : public XlaOpKernel {
     //                            [0, 0, 0, 4]]
 
     // Flattens the input to 1D.
+    xla::DExpr flattened_expr = xla::DExpr::Const(1);
+    std::vector<xla::DExpr> input_exprs = input_shape.get_filled_expressions();
+    for (const xla::DExpr& expr : input_exprs) {
+      flattened_expr = (flattened_expr * expr).simplify();
+    }
     int64_t size = input_shape.num_elements();
-    input = xla::Reshape(input, {size});
+    input = xla::Reshape(input, {size}, {flattened_expr});
 
     // Create an R2 with the R1 diagonal.
-    xla::XlaOp diag = CreateDiagonal(input, size, /*other_dims=*/{});
+    xla::XlaOp diag =
+        CreateDiagonal(input, size, flattened_expr, /*other_dims=*/{},
+                       /*other_dim_exprs=*/{});
 
     // Reshapes to the final shape.
     std::vector<int64_t> new_dims(dims.size() * 2);
     std::copy(dims.begin(), dims.end(), new_dims.begin());
     std::copy(dims.begin(), dims.end(), new_dims.begin() + dims.size());
-    diag = xla::Reshape(diag, new_dims);
+    std::vector<xla::DExpr> new_exprs(input_exprs.begin(), input_exprs.end());
+    new_exprs.insert(new_exprs.end(), input_exprs.begin(), input_exprs.end());
+    diag = xla::Reshape(diag, new_dims, new_exprs);
 
     ctx->SetOutput(0, diag);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc
index ceeea010ee7858..46404dc31ed23f 100644
--- a/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/dynamic_partition_op.cc
@@ -43,6 +43,26 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+std::vector<xla::DExpr> GetFilledExpressions(const xla::Shape& shape) {
+  std::vector<xla::DExpr> expressions;
+  expressions.reserve(shape.dimensions_size());
+  auto shape_expressions = shape.expressions();
+  for (int64_t i = 0; i < shape.dimensions_size(); ++i) {
+    expressions.push_back(i < shape_expressions.size() && shape_expressions[i]
+                              ? shape_expressions[i]
+                              : xla::DExpr::Const(shape.dimensions(i)));
+  }
+  return expressions;
+}
+
+xla::DExpr CollapseExpressions(absl::Span<const xla::DExpr> expressions) {
+  xla::DExpr collapsed = xla::DExpr::Const(1);
+  for (const xla::DExpr& expression : expressions) {
+    collapsed = collapsed * expression;
+  }
+  return collapsed;
+}
+
 class DynamicPartitionOp : public XlaOpKernel {
  public:
   explicit DynamicPartitionOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
@@ -67,6 +87,26 @@ class DynamicPartitionOp : public XlaOpKernel {
                      xla::XlaOp partitions_1d, const xla::Shape& data_1d_shape,
                      const xla::Shape& partition_1d_shape) {
     int64_t input_count = data_1d_shape.dimensions(0);
+    VLOG(1) << "data_1d_shape="
+          << xla::ShapeUtil::HumanString(data_1d_shape)
+          << " partition_1d_shape="
+          << xla::ShapeUtil::HumanString(partition_1d_shape)
+          << " input_count=" << input_count
+          << " num_partitions=" << num_partitions_;
+    // Use the smaller runtime size to mask off padded tail elements when one
+    // input loses dynamic shape information and falls back to its static bound.
+    // For valid DynamicPartition inputs the true prefix sizes should match, so
+    // this keeps us from treating padding on either side as real elements.
+    xla::XlaOp dynamic_input_count = xla::Min(
+        xla::GetDimensionSize(data_1d, 0),
+        xla::GetDimensionSize(partitions_1d, 0));
+    xla::XlaOp input_index = xla::Iota(ctx->builder(), xla::S32, input_count);
+    xla::XlaOp valid_element = xla::Lt(input_index, dynamic_input_count);
+    xla::XlaOp invalid_partition =
+        xla::Broadcast(xla::ConstantR0<int32>(ctx->builder(), num_partitions_),
+                       {input_count}, partition_1d_shape.expressions());
+    partitions_1d = xla::Select(valid_element, partitions_1d, invalid_partition);
+
     std::vector<xla::XlaOp> to_sort = {partitions_1d, data_1d};
     std::vector<xla::PrimitiveType> types_to_sort = {
         partition_1d_shape.element_type(), data_1d_shape.element_type()};
@@ -114,11 +154,13 @@ class DynamicPartitionOp : public XlaOpKernel {
   void Compile(XlaOpKernelContext* ctx) override {
     xla::Shape data_shape = ctx->InputXlaShape(0).value();
     xla::Shape partition_shape = ctx->InputXlaShape(1).value();
+    xla::Shape flattened_partition_shape = partition_shape;
     xla::XlaOp data = ctx->Input(0);
     xla::XlaOp partitions = ctx->Input(1);
     std::vector<int64_t> partitions_static;
     bool partitions_are_static =
-        ctx->ConstantInputReshapedToIntVector(1, &partitions_static).ok();
+      ctx->ConstantInputReshapedToIntVector(1, &partitions_static).ok() &&
+      partition_shape.is_static();
     // We know how to solve DynamicPartition on 1D inputs using
     // DynamicPartition1D. For other input, we do two things:
     //
@@ -140,8 +182,11 @@ class DynamicPartitionOp : public XlaOpKernel {
       for (int64_t i = 0; i < rank; ++i) {
         broadcasted_dims.push_back(i);
       }
-      partitions = xla::BroadcastInDim(partitions, data_shape.dimensions(),
-                                       broadcasted_dims);
+      partitions =
+          xla::BroadcastInDim(partitions, data_shape.dimensions(),
+                              broadcasted_dims, data_shape.expressions());
+      flattened_partition_shape = data_shape;
+      flattened_partition_shape.set_element_type(partition_shape.element_type());
     }
 
     // Output shape bounded is calculated by
@@ -159,19 +204,39 @@ class DynamicPartitionOp : public XlaOpKernel {
     }
 
     int64_t input_count = xla::ShapeUtil::ElementsIn(data_shape);
-    auto data_1d = xla::Reshape(data, {input_count});
-    auto partitions_1d = xla::Reshape(partitions, {input_count});
-    xla::Shape data_1d_shape =
-        xla::ShapeUtil::MakeShape(data_shape.element_type(), {input_count});
+    auto data_exprs = GetFilledExpressions(data_shape);
+    auto flattened_partition_exprs = GetFilledExpressions(flattened_partition_shape);
+    auto data_1d =
+      xla::Reshape(data, {input_count}, {CollapseExpressions(data_exprs)});
+    auto partitions_1d = xla::Reshape(
+      partitions, {input_count},
+      {CollapseExpressions(flattened_partition_exprs)});
+    xla::Shape data_1d_shape = xla::ShapeUtil::MakeShape(
+        data_shape.element_type(), {input_count},
+        std::vector<xla::DExpr>{CollapseExpressions(data_exprs)});
 
     xla::Shape partitions_1d_shape = xla::ShapeUtil::MakeShape(
-        partition_shape.element_type(), {input_count});
+        partition_shape.element_type(), {input_count},
+        std::vector<xla::DExpr>{
+            CollapseExpressions(flattened_partition_exprs)});
 
     std::vector<xla::XlaOp> output, partition_length;
     std::tie(output, partition_length) = DynamicPartition1D(
         ctx, data_1d, partitions_1d, data_1d_shape, partitions_1d_shape);
+    std::vector<xla::DExpr> output_exprs;
+    output_exprs.reserve(data_shape.dimensions().size() -
+                         partition_shape.dimensions().size() + 1);
+    output_exprs.push_back(CollapseExpressions(absl::MakeConstSpan(
+      flattened_partition_exprs).subspan(
+      0, partition_shape.dimensions().size())));
+    for (int64_t i = partition_shape.dimensions().size();
+         i < data_shape.dimensions().size(); ++i) {
+      output_exprs.push_back(data_exprs[i]);
+    }
     for (int64_t i = 0; i < num_partitions_; ++i) {
-      auto reshape = xla::Reshape(output[i], output_shape_bound_dims);
+      xla::Shape output_shape = xla::ShapeUtil::MakeShape(
+          data_shape.element_type(), output_shape_bound_dims, output_exprs);
+      auto reshape = xla::Reshape(output_shape, output[i]);
       if (partitions_are_static) {
         int64_t size = absl::c_count(partitions_static, i);
         ctx->SetOutput(i, xla::SliceInDim(reshape, 0, size, 1, 0));
diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
index cb7e4f6f96437e..fe45493d452d46 100644
--- a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc
@@ -132,13 +132,18 @@ class DynamicStitchOp : public XlaOpKernel {
     int64_t result_rank = 1 + data0_shape.dims() - indices0_shape.dims();
     if (number_of_indices == 0) {
       std::vector<int64_t> result_shape(result_rank);
+      std::vector<xla::DExpr> result_expressions(result_rank,
+                                                 xla::DExpr::Const(0));
       for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) {
         result_shape[d - indices0_shape.dims() + 1] = data0_shape.dim_size(d);
+        result_expressions[d - indices0_shape.dims() + 1] =
+            data0_shape.get_filled_expression(d);
       }
       xla::PrimitiveType element_type =
           ctx->input_xla_type(ctx->num_inputs() - 1);
       xla::Literal empty_literal = xla::Literal::CreateFromShape(
-          xla::ShapeUtil::MakeShape(element_type, result_shape));
+          xla::ShapeUtil::MakeShape(element_type, result_shape,
+                                    result_expressions));
       ctx->SetOutput(0, xla::ConstantLiteral(ctx->builder(), empty_literal));
       return;
     }
@@ -174,16 +179,20 @@ class DynamicStitchOp : public XlaOpKernel {
       TensorShape new_shape;
       // first reshaped dimension is the number of indices for this input.
       new_shape.AddDim(indices[input_num].shape().dimensions(0));
+      new_shape.AddExpression(
+          xla::DExpr::Const(indices[input_num].shape().dimensions(0)));
       // Then the rest are the common extra shape.
       for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) {
         new_shape.AddDim(data0_shape.dim_size(d));
+        new_shape.AddExpression(data0_shape.get_filled_expression(d));
       }
       // Get the data, shaped appropriately.
       auto handle = data[input_num];
       if (new_shape == data_shapes[input_num]) {
         input[input_num] = handle;
       } else {
-        input[input_num] = xla::Reshape(handle, new_shape.dim_sizes());
+        input[input_num] = xla::Reshape(handle, new_shape.dim_sizes(),
+                                        new_shape.get_filled_expressions());
       }
     }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc
index 2a65441eb79bf9..e7bc0890317316 100644
--- a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc
@@ -127,7 +127,8 @@ class FakeQuantWithMinMaxArgsGradOp : public XlaOpKernel {
     xla::XlaOp between_nudged_min_max = xla::And(
         xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max));
     xla::XlaOp zeroes = xla::Broadcast(XlaHelpers::Zero(b, data_type),
-                                       gradient_shape.dim_sizes());
+                                       gradient_shape.dim_sizes(),
+                                       gradient_shape.get_filled_expressions());
     xla::XlaOp output = xla::Select(between_nudged_min_max, gradient, zeroes);
     ctx->SetOutput(0, output);
   }
@@ -213,7 +214,8 @@ class FakeQuantWithMinMaxVarsGradOp : public XlaOpKernel {
     xla::XlaOp between_nudged_min_max = xla::And(
         xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max));
     xla::XlaOp zero = XlaHelpers::Zero(b, data_type);
-    xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes());
+    xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes(),
+                                       gradient_shape.get_filled_expressions());
     xla::XlaOp output0 = xla::Select(between_nudged_min_max, gradient, zeroes);
     ctx->SetOutput(0, output0);
 
@@ -270,7 +272,8 @@ class FakeQuantWithMinMaxVarsPerChannelOp : public XlaOpKernel {
     absl::Span<const int64_t> input_dimensions = input_shape.dimensions();
     auto convert_to_input_shape = [&](const xla::XlaOp op) {
       return xla::BroadcastInDim(op, input_dimensions,
-                                 {input_shape.dimensions_size() - 1});
+                                 {input_shape.dimensions_size() - 1},
+                                 input_shape.expressions());
     };
     input_min = convert_to_input_shape(input_min);
     input_max = convert_to_input_shape(input_max);
@@ -323,7 +326,6 @@ class FakeQuantWithMinMaxVarsPerChannelGradOp : public XlaOpKernel {
     xla::XlaBuilder* b = ctx->builder();
     xla::Shape input_shape = b->GetShape(input).value();
     absl::Span<const int64_t> input_dimensions = input_shape.dimensions();
-
     std::vector<int64_t> reduce_axes;
     for (int64_t i = 0; i + 1 < input_shape.dimensions_size(); ++i) {
       reduce_axes.push_back(i);
@@ -331,7 +333,8 @@ class FakeQuantWithMinMaxVarsPerChannelGradOp : public XlaOpKernel {
 
     auto convert_to_input_shape = [&](const xla::XlaOp op) {
       return xla::BroadcastInDim(op, input_dimensions,
-                                 {input_shape.dimensions_size() - 1});
+                                 {input_shape.dimensions_size() - 1},
+                                 input_shape.expressions());
     };
     input_min = convert_to_input_shape(input_min);
     input_max = convert_to_input_shape(input_max);
@@ -343,7 +346,8 @@ class FakeQuantWithMinMaxVarsPerChannelGradOp : public XlaOpKernel {
     xla::XlaOp between_nudged_min_max = xla::And(
         xla::Le(nudged_input_min, input), xla::Le(input, nudged_input_max));
     xla::XlaOp zero = XlaHelpers::Zero(b, data_type);
-    xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes());
+    xla::XlaOp zeroes = xla::Broadcast(zero, gradient_shape.dim_sizes(),
+                                       gradient_shape.get_filled_expressions());
     xla::XlaOp output0 = xla::Select(between_nudged_min_max, gradient, zeroes);
     ctx->SetOutput(0, output0);
 
diff --git a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
index 8fb04773aafb49..5d4114d8c88d5d 100644
--- a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc
@@ -20,8 +20,8 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/inlined_vector.h"
-#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/xla_builder.h"
@@ -139,9 +139,9 @@ class IFFTOp : public GenericFftOp {
   explicit IFFTOp(OpKernelConstruction* ctx)
       : GenericFftOp(ctx, /*fft_type=*/FftType::IFFT, /*fft_rank=*/FFTRank) {}
 };
-REGISTER_XLA_OP(Name("IFFT").TypeConstraint("Tcomplex",
-                                            {DT_COMPLEX64, DT_COMPLEX128}),
-                MlirXlaOpKernel);
+REGISTER_XLA_OP_FACTORY(
+    Name("IFFT").TypeConstraint("Tcomplex", {DT_COMPLEX64, DT_COMPLEX128}),
+    CreateDynamicNativeXlaOpKernel<IFFTOp<1>>);
 REGISTER_XLA_OP(Name("IFFT2D").TypeConstraint("Tcomplex",
                                               {DT_COMPLEX64, DT_COMPLEX128}),
                 IFFTOp<2>);
diff --git a/tensorflow/compiler/tf2xla/kernels/fill_op.cc b/tensorflow/compiler/tf2xla/kernels/fill_op.cc
index 6e5a1430538365..97e4c6e85d4dcf 100644
--- a/tensorflow/compiler/tf2xla/kernels/fill_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/fill_op.cc
@@ -53,11 +53,20 @@ class FillOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx,
                    ctx->ConstantInputAsIntVector(
                        "dims", &dims, xla::ValueInferenceMode::kUpperBound));
+    std::vector<xla::DExpr> dim_exprs;
+    dim_exprs.reserve(dims.size());
+    const auto& contents = ctx->InputExpression("dims").contents();
+    for (int64_t i = 0; i < dims.size(); ++i) {
+      dim_exprs.push_back(i < contents.size() && contents[i] &&
+                                  contents[i]->is_dynamic()
+                              ? contents[i]
+                              : xla::DExpr::Const(dims[i]));
+    }
     std::vector<bool> dynamic_dims;
     OP_REQUIRES_OK(
         ctx, ctx->ResolveInputDynamismIntoPredVector("dims", &dynamic_dims));
 
-    auto output = xla::Broadcast(ctx->Input("value"), dims);
+    auto output = xla::Broadcast(ctx->Input("value"), dims, dim_exprs);
     for (int64_t i = 0; i < dims.size(); ++i) {
       // If a dimension is dynamic, call set-dimension-size on the output.
       if (dynamic_dims[i]) {
diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc
index 2783951e1b6b0f..f058f5a9722506 100644
--- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc
@@ -88,7 +88,8 @@ absl::Status XlaGather(const xla::XlaOp& input, const TensorShape& input_shape,
     out_shape.AppendShape(input_shape_post_axis);
 
     *gather_output =
-        xla::Broadcast(XlaHelpers::Zero(builder, dtype), out_shape.dim_sizes());
+        xla::Broadcast(XlaHelpers::Zero(builder, dtype), out_shape.dim_sizes(),
+                       out_shape.get_filled_expressions());
     return absl::OkStatus();
   }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/identity_op.cc b/tensorflow/compiler/tf2xla/kernels/identity_op.cc
index 3e765b853e110d..f6abb610c456e4 100644
--- a/tensorflow/compiler/tf2xla/kernels/identity_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/identity_op.cc
@@ -32,6 +32,10 @@ class IdentityOp : public XlaOpKernel {
     for (int i = 0; i < ctx->num_inputs(); ++i) {
       if (IsTensorListInput(ctx, i)) {
         ctx->SetTensorListOutput(i, ctx->Input(i));
+      } else if (ctx->InputExpression(i).kind() !=
+                     XlaExpression::Kind::kResource &&
+                 ctx->input_type(i) != DT_VARIANT) {
+        ctx->SetOutputExpression(i, ctx->InputExpression(i));
       } else {
         DCHECK(ctx->input_type(i) != DT_VARIANT);
         // Forwards using the underlying op_kernel_context so both tensor and
@@ -58,7 +62,8 @@ REGISTER_XLA_OP(Name("IdentityN")
                     .CompilationOnly(),
                 IdentityOp);
 REGISTER_XLA_OP(Name("PlaceholderWithDefault"), IdentityOp);
-REGISTER_XLA_OP(Name("PreventGradient"), MlirXlaOpKernel);
+REGISTER_XLA_OP_FACTORY(
+    Name("PreventGradient"), CreateDynamicNativeXlaOpKernel<IdentityOp>);
 REGISTER_XLA_OP(Name("StopGradient").AllowVariantTypes(), IdentityOp);
 REGISTER_XLA_OP(Name("Snapshot"), IdentityOp);
 REGISTER_XLA_OP(Name("_EagerConst"), IdentityOp);
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index a8eb7bbf794268..6cc0b873a355fd 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -59,7 +59,7 @@ std::array<xla::XlaOp, 3> RGBToHSV(XlaOpKernelContext* ctx, xla::XlaBuilder* b,
   auto minimum = xla::Min(xla::Min(red, green), blue);
   auto range = xla::Sub(value, minimum);
 
-  auto zeros = xla::Broadcast(zero, shape.dim_sizes());
+  auto zeros = xla::Broadcast(zero, shape.dim_sizes(), shape.get_filled_expressions());
   auto saturation =
       xla::Select(xla::Gt(value, zero), xla::Div(range, value), zeros);
 
diff --git a/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc b/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc
index f357262a39c35b..3808df4937b3ec 100644
--- a/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc
@@ -87,9 +87,11 @@ class InTopKOp : public XlaOpKernel {
     // which indicates the target is in topk.
     xla::XlaOp gt_r2 = xla::Gt(predictions_r2, targets_values_r1, {0});
     xla::XlaOp zero_r0 = xla::Zero(xla_builder, xla::S32);
-    xla::XlaOp zero_r2 = xla::Broadcast(zero_r0, predictions_shape.dim_sizes());
+    xla::XlaOp zero_r2 = xla::Broadcast(zero_r0, predictions_shape.dim_sizes(),
+                                        predictions_shape.get_filled_expressions());
     xla::XlaOp one_r0 = xla::One(xla_builder, xla::S32);
-    xla::XlaOp one_r2 = xla::Broadcast(one_r0, predictions_shape.dim_sizes());
+    xla::XlaOp one_r2 = xla::Broadcast(one_r0, predictions_shape.dim_sizes(),
+                                       predictions_shape.get_filled_expressions());
     xla::XlaOp one_hot_r2 = xla::Select(gt_r2, one_r2, zero_r2);
     xla::XlaOp num_gt_r1 = xla::Reduce(
         one_hot_r2, zero_r0,
diff --git a/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc b/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc
index 46e46f6d8b3d32..cf2389110eefb4 100644
--- a/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc
@@ -49,14 +49,16 @@ void BuildLowerUpperBoundOp(XlaOpKernelContext* ctx, DataType out_dtype,
   // dimension of sorted_sequence.
   auto new_values_shape = values_shape;
   new_values_shape.InsertDim(/* d */ 2, /* size */ 1);
-  auto values_reshaped = xla::Reshape(values, new_values_shape.dim_sizes());
+  auto values_reshaped = xla::Reshape(values, new_values_shape.dim_sizes(),
+                                      new_values_shape.get_filled_expressions());
 
   // Add a new penultimate dimension to sorted_inputs, to allow broadcasting of
   // sorted_sequence entries for each value.
   auto new_sorted_inputs_shape = sorted_inputs_shape;
   new_sorted_inputs_shape.InsertDim(/* d */ 1, /* size */ 1);
   auto sorted_inputs_reshaped =
-      xla::Reshape(sorted_inputs, new_sorted_inputs_shape.dim_sizes());
+      xla::Reshape(sorted_inputs, new_sorted_inputs_shape.dim_sizes(),
+                   new_sorted_inputs_shape.get_filled_expressions());
 
   // We are relying on broadcasting to compare each value against each entry in
   // the associated sorted_inputs row.
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
index 48e8f976cc67bb..8ef6dfbd9a1702 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
@@ -234,7 +234,8 @@ xla::XlaOp SetMatrixDiag(const xla::XlaOp input, const xla::XlaOp diag,
 
     // Broadcast and mask.
     xla::XlaOp diag_broadcast = xla::BroadcastInDim(
-        diag_slice, input_shape.dim_sizes(), broadcast_dimensions);
+        diag_slice, input_shape.dim_sizes(), broadcast_dimensions,
+        input_shape.get_filled_expressions());
     const auto mask = xla::GetDiagonalMask(output, diag_index);
     output = xla::Select(mask, diag_broadcast, output);
   }
@@ -327,8 +328,11 @@ class MatrixDiagOp : public XlaOpKernel {
     TensorShape output_shape = diag_shape;
     output_shape.RemoveLastDims((num_diags == 1) ? 1 : 2);
     output_shape.AddDim(num_rows);
+    output_shape.AddExpression(xla::DExpr::Const(num_rows));
     output_shape.AddDim(num_cols);
-    xla::XlaOp output = xla::Broadcast(padding_value, output_shape.dim_sizes());
+    output_shape.AddExpression(xla::DExpr::Const(num_cols));
+    xla::XlaOp output = xla::Broadcast(padding_value, output_shape.dim_sizes(),
+                                       output_shape.get_filled_expressions());
     xla::XlaOp diag = context->Input(0);
     context->SetOutput(
         0, SetMatrixDiag(output, diag, output_shape, diag_rank, num_diags,
@@ -404,11 +408,15 @@ class MatrixDiagPartOp : public XlaOpKernel {
     TensorShape output_shape = input_shape;
     output_shape.RemoveLastDims(2);
     const int num_diags = upper_diag_index - lower_diag_index + 1;
-    if (num_diags > 1) output_shape.AddDim(num_diags);
+    if (num_diags > 1) {
+      output_shape.AddDim(num_diags);
+      output_shape.AddExpression(xla::DExpr::Const(num_diags));
+    }
     const int32_t max_diag_len =
         std::min(num_rows + std::min(upper_diag_index, int64_t{0}),
                  num_cols - std::max(lower_diag_index, int64_t{0}));
     output_shape.AddDim(max_diag_len);
+    output_shape.AddExpression(xla::DExpr::Const(max_diag_len));
 
     // Computes output.
     xla::XlaOp input = context->Input(0);
@@ -447,7 +455,8 @@ class MatrixDiagPartOp : public XlaOpKernel {
     }
     auto concat =
         xla::ConcatInDim(context->builder(), diag_list, input_rank - 2);
-    context->SetOutput(0, xla::Reshape(concat, output_shape.dim_sizes()));
+    context->SetOutput(0, xla::Reshape(concat, output_shape.dim_sizes(),
+                                       output_shape.get_filled_expressions()));
   }
 
  private:
@@ -519,11 +528,15 @@ class MatrixSetDiagOp : public XlaOpKernel {
 
     TensorShape expected_diag_shape = input_shape;
     expected_diag_shape.RemoveLastDims(2);
-    if (num_diags > 1) expected_diag_shape.AddDim(num_diags);
+    if (num_diags > 1) {
+      expected_diag_shape.AddDim(num_diags);
+      expected_diag_shape.AddExpression(xla::DExpr::Const(num_diags));
+    }
     const int32_t max_diag_len =
         std::min(num_rows + std::min(upper_diag_index, int64_t{0}),
                  num_cols - std::max(lower_diag_index, int64_t{0}));
     expected_diag_shape.AddDim(max_diag_len);
+    expected_diag_shape.AddExpression(xla::DExpr::Const(max_diag_len));
     OP_REQUIRES(
         context, expected_diag_shape == diag_shape,
         errors::InvalidArgument(
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
index 17b5ae7a70375a..207bc6f7ce8984 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc
@@ -96,8 +96,12 @@ MatrixTriangularSolveOp::Broadcast(xla::XlaOp lhs, const TensorShape& lhs_shape,
 
   TensorShape lhs_broadcast_shape(broadcast_helper.output_batch_shape());
   lhs_broadcast_shape.AddDim(m);
+  lhs_broadcast_shape.AddExpression(xla::DExpr::Const(m));
   lhs_broadcast_shape.AddDim(m);
-  auto lhs_output = BroadcastTo(lhs, lhs_broadcast_shape.dim_sizes());
+  lhs_broadcast_shape.AddExpression(xla::DExpr::Const(m));
+  auto lhs_output =
+      BroadcastTo(lhs, lhs_broadcast_shape.dim_sizes(),
+                  lhs_broadcast_shape.get_filled_expressions());
   if (!lhs_output.ok()) {
     xla::XlaOp error = lhs.builder()->ReportError(lhs_output.status());
     return {error, error};
@@ -105,8 +109,12 @@ MatrixTriangularSolveOp::Broadcast(xla::XlaOp lhs, const TensorShape& lhs_shape,
 
   TensorShape rhs_broadcast_shape(broadcast_helper.output_batch_shape());
   rhs_broadcast_shape.AddDim(m);
+  rhs_broadcast_shape.AddExpression(xla::DExpr::Const(m));
   rhs_broadcast_shape.AddDim(n);
-  auto rhs_output = BroadcastTo(rhs, rhs_broadcast_shape.dim_sizes());
+  rhs_broadcast_shape.AddExpression(xla::DExpr::Const(n));
+  auto rhs_output =
+      BroadcastTo(rhs, rhs_broadcast_shape.dim_sizes(),
+                  rhs_broadcast_shape.get_filled_expressions());
   if (!rhs_output.ok()) {
     xla::XlaOp error = rhs.builder()->ReportError(rhs_output.status());
     return {error, error};
diff --git a/tensorflow/compiler/tf2xla/kernels/pack_op.cc b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
index ba4e8bbef7b136..0b236ae9c0c052 100644
--- a/tensorflow/compiler/tf2xla/kernels/pack_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
@@ -17,6 +17,8 @@ limitations under the License.
 
 #include <vector>
 
+#include "absl/algorithm/container.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/xla_builder.h"
@@ -28,6 +30,33 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+bool TryBuildPackedContents(XlaOpKernelContext* ctx, int num, int axis,
+                            std::vector<xla::DExpr>* contents) {
+  contents->clear();
+  if (axis != 0) {
+    return false;
+  }
+  for (int i = 0; i < num; ++i) {
+    if (ctx->InputShape(i).dims() != 0) {
+      contents->clear();
+      return false;
+    }
+    const auto& input_contents = ctx->InputExpression(i).contents();
+    if (!input_contents.empty()) {
+      if (input_contents.size() != 1) {
+        contents->clear();
+        return false;
+      }
+      contents->push_back(input_contents[0]);
+      continue;
+    }
+    contents->push_back(xla::DExpr::Unknown(xla::kUnknownContentSentinel));
+  }
+  return absl::c_any_of(*contents, [](const xla::DExpr& expr) {
+    return expr && expr->is_dynamic();
+  });
+}
+
 class PackOp : public XlaOpKernel {
  public:
   explicit PackOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
@@ -64,14 +93,36 @@ class PackOp : public XlaOpKernel {
     std::vector<xla::XlaOp> reshaped_inputs(num);
 
     TensorShape child_shape(shapes[0]);
+    std::vector<xla::DExpr> exprs = child_shape.get_filled_expressions();
     child_shape.InsertDim(axis, 1);
-
+    exprs.insert(exprs.begin() + axis, xla::DExpr::Const(1));
+    std::vector<xla::DExpr> output_contents;
+    const bool has_output_contents =
+        SymbolicContentEnabled() &&
+        TryBuildPackedContents(ctx, num, axis, &output_contents);
     for (int i = 0; i < num; ++i) {
       // Reshape the inputs to have an extra dimension of size 1.
-      reshaped_inputs[i] = xla::Reshape(values[i], child_shape.dim_sizes());
+      reshaped_inputs[i] = xla::Reshape(values[i], child_shape.dim_sizes(),
+                                        exprs);
+      if (has_output_contents) {
+        OP_REQUIRES_OK(
+            ctx, ctx->builder()->SetInstructionContents(
+                     reshaped_inputs[i],
+                     {output_contents[static_cast<size_t>(i)]}));
+      }
     }
 
-    ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), reshaped_inputs, axis));
+    auto output = xla::ConcatInDim(ctx->builder(), reshaped_inputs, axis);
+    if (has_output_contents) {
+      OP_REQUIRES_OK(
+          ctx, ctx->builder()->SetInstructionContents(output, output_contents));
+      auto output_expr =
+          XlaExpression::XlaOp(output, ctx->expected_output_dtype(0));
+      output_expr.set_contents(std::move(output_contents));
+      ctx->SetOutputExpression(0, output_expr);
+      return;
+    }
+    ctx->SetOutput(0, output);
   }
 
  private:
diff --git a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
index aa7c78b8b8f97a..fe2401fade704a 100644
--- a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc
@@ -264,10 +264,14 @@ class MaxPoolOp : public PoolingOp {
 
       absl::InlinedVector<int64, 5> new_dims(result_shape->dimensions().begin(),
                                              result_shape->dimensions().end());
+      absl::InlinedVector<xla::DExpr, 5> new_exprs(
+          result_shape->expressions().begin(),
+          result_shape->expressions().end());
       new_dims[1] /= *vect_width;
+      new_exprs[1] = new_exprs[1] / xla::DExpr::Const(*vect_width);
       new_dims.insert(new_dims.begin() + 2, *vect_width);
-      pooling =
-          xla::Transpose(xla::Reshape(pooling, new_dims), {0, 1, 3, 4, 2});
+      pooling = xla::Transpose(xla::Reshape(pooling, new_dims, new_exprs),
+                               {0, 1, 3, 4, 2});
     }
 
     ctx->SetOutput(0, pooling);
diff --git a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
index cac9f8a68f234e..6b7e990f7030fd 100644
--- a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
@@ -174,7 +174,8 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
       xla::Shape input_shape = b->GetShape(input).value();
       absl::Span<const int64_t> input_dimensions = input_shape.dimensions();
       auto convert_to_input_shape = [&](const xla::XlaOp op) {
-        return xla::BroadcastInDim(op, input_dimensions, {axis_});
+        return xla::BroadcastInDim(op, input_dimensions, {axis_},
+                                   input_shape.expressions());
       };
       min_range = convert_to_input_shape(min_range);
       max_range = convert_to_input_shape(max_range);
diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc
index 5f911018c244b5..f1342c7e187ea2 100644
--- a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/constants.h"
@@ -147,10 +148,10 @@ class MeanOp : public XlaReductionOp {
     xla::XlaOp result = reduce_output;
     xla::Shape bounded_shape = builder->GetShape(input).value();
     int64_t divisor_value = bounded_shape.dimensions(dimensions_to_reduce[0]);
-    auto divisor = xla::GetDimensionSize(input, dimensions_to_reduce[0]);
+    xla::XlaOp divisor = xla::GetDimensionSize(input, dimensions_to_reduce[0]);
     for (int i = 1; i < dimensions_to_reduce.size(); i++) {
       int64_t size_value = bounded_shape.dimensions(dimensions_to_reduce[i]);
-      auto size = xla::GetDimensionSize(input, dimensions_to_reduce[i]);
+      xla::XlaOp size = xla::GetDimensionSize(input, dimensions_to_reduce[i]);
       if (size_value * divisor_value > std::numeric_limits<int32_t>::max()) {
         result = result / xla::ConvertElementType(divisor, xla_reduction_type_);
         divisor_value = size_value;
diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc
index 6a8a98342c1123..6a5f40441beff3 100644
--- a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc
@@ -106,16 +106,19 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) {
   }
 
   std::vector<int64_t> final_shape;
+  std::vector<xla::DExpr> final_exprs;
   for (int i = 0; i < data_shape.dims(); ++i) {
     if (!bitmap[i]) {
       // If we are not reducing along dimension i.
       int64_t dim = data_shape.dim_size(i);
       final_shape.push_back(dim);
+      final_exprs.push_back(data_shape.get_filled_expression(i));
     } else if (keep_dims_) {
       // We are reducing along dimension i, but we want to keep the
       // same number of dimensions, so we set the dimension of i to
       // '1'.
       final_shape.push_back(1);
+      final_exprs.push_back(xla::DExpr::Const(1));
     }
   }
 
@@ -139,7 +142,8 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) {
 
   auto reduce = xla::Reduce(data, initial, reduction_computation, xla_axes);
   auto finalized = BuildFinalizer(b, data, reduce, xla_axes);
-  auto result = keep_dims_ ? xla::Reshape(finalized, final_shape) : finalized;
+  auto result = keep_dims_ ? xla::Reshape(finalized, final_shape, final_exprs)
+                           : finalized;
   ctx->SetOutput(0, result);
 }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/relu_op.cc b/tensorflow/compiler/tf2xla/kernels/relu_op.cc
index f274b271596ff5..fd932247a0f558 100644
--- a/tensorflow/compiler/tf2xla/kernels/relu_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/relu_op.cc
@@ -40,8 +40,28 @@ XlaOp Relu6(XlaOp x) {
 namespace tensorflow {
 namespace {
 
-REGISTER_XLA_OP(Name("Relu"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("Relu6"), MlirXlaOpKernel);
+class ReluOp : public XlaOpKernel {
+ public:
+  explicit ReluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  // Computes the max of the scalar input x and 0.
+  void Compile(XlaOpKernelContext* ctx) override {
+    ctx->SetOutput(0, xla::Relu(ctx->Input(0)));
+  }
+};
+
+class Relu6Op : public XlaOpKernel {
+ public:
+  explicit Relu6Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  // Computes the max of the scalar input x and 0.
+  void Compile(XlaOpKernelContext* ctx) override {
+    ctx->SetOutput(0, xla::Relu6(ctx->Input(0)));
+  }
+};
+
+REGISTER_XLA_OP(Name("Relu"), ReluOp);
+// REGISTER_XLA_OP(Name("Relu"), MlirXlaOpKernel);
+// REGISTER_XLA_OP(Name("Relu6"), MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("Relu6"), Relu6Op);
 
 class LeakyReluOp : public XlaOpKernel {
  public:
@@ -70,9 +90,11 @@ class Relu6GradOp : public XlaOpKernel {
     xla::XlaBuilder* b = ctx->builder();
     const TensorShape shape = ctx->InputShape(0);
     const auto zero =
-        xla::Broadcast(XlaHelpers::Zero(b, input_type(0)), shape.dim_sizes());
-    const auto six = xla::Broadcast(
-        XlaHelpers::IntegerLiteral(b, input_type(0), 6), shape.dim_sizes());
+        xla::Broadcast(XlaHelpers::Zero(b, input_type(0)), shape.dim_sizes(),
+                       shape.get_filled_expressions());
+    const auto six =
+        xla::Broadcast(XlaHelpers::IntegerLiteral(b, input_type(0), 6),
+                       shape.dim_sizes(), shape.get_filled_expressions());
     auto out = xla::Select(
         xla::And(xla::Lt(ctx->Input(1), six), xla::Gt(ctx->Input(1), zero)),
         ctx->Input(0), zero);
diff --git a/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc
index c54c4613d29e44..602d480b5b6541 100644
--- a/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc
@@ -122,12 +122,15 @@ XlaOp ConcatenateIota(xla::XlaBuilder* b, XlaOp indices,
   for (auto dim : warp_shape) {
     dimensions.push_back(dim.size);
   }
+  std::vector<xla::DExpr> expressions =
+      warp_shape.get_filled_expressions();
   // Except the last dimension, which is of size 1.
   dimensions.back() = 1;
+  expressions.back() = xla::DExpr::Const(1);
 
-  auto batch_indices =
-      xla::Iota(b, xla::ShapeUtil::MakeShape(xla::S32, dimensions),
-                /*iota_dimension=*/0);
+  auto batch_indices = xla::Iota(
+      b, xla::ShapeUtil::MakeShape(xla::S32, dimensions, expressions),
+      /*iota_dimension=*/0);
 
   return xla::ConcatInDim(b, {batch_indices, indices}, dimensions.size() - 1);
 }
@@ -365,14 +368,17 @@ XlaOp CalculateGradWarp(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio,
   auto warp_dims = warp_shape.dim_sizes();
   std::vector<int64_t> warp_dims_without_last_dims(warp_dims.begin(),
                                                    warp_dims.end() - 1);
+  std::vector<xla::DExpr> warp_expressions =
+      warp_shape.get_filled_expressions();
 
   // With dimension [batch, dim_0, ...dim_n, 4]
   std::vector<int64_t> neighbor_broadcast_dims = warp_dims_without_last_dims;
   neighbor_broadcast_dims.push_back(4);
+  warp_expressions.back() = xla::DExpr::Const(4);
 
   // With dimension [batch, dim_0, ...dim_n, 4]
-  auto neighbor_broadcast_shape =
-      xla::ShapeUtil::MakeShape(data_type, neighbor_broadcast_dims);
+  auto neighbor_broadcast_shape = xla::ShapeUtil::MakeShape(
+      data_type, neighbor_broadcast_dims, warp_expressions);
 
   const int64_t last_warp_dim = warp_shape.dims() - 1;
 
diff --git a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc
index ba17d1b295b763..c019f42927bb9b 100644
--- a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc
@@ -57,8 +57,10 @@ class ReshapeOp : public XlaOpKernel {
     // is one.
     TensorShape shape;
     int64_t product = 1;
+    xla::DExpr product_expr = xla::DExpr::Const(1);
     int unknown_index = -1;
     bool shape_has_zero_dim = false;
+    int ratio = 1;
     for (int d = 0; d < num_dims; ++d) {
       const int64_t size = shape_input[d];
       if (size == -1) {
@@ -68,23 +70,60 @@ class ReshapeOp : public XlaOpKernel {
                                     unknown_index, " and ", d));
         unknown_index = d;
         shape.AddDim(1);
+        shape.AddExpression(xla::DExpr::Const(1));
+        ratio = 1;
       } else if (size == 0) {
         // We don't include zero-sized dimension in product, so that we can
         // still calculate number of elements for non-zero-sized dimensions and
         // therefore infer their shapes.
         shape.AddDim(size);
+        shape.AddExpression(xla::DExpr::Const(size));
         shape_has_zero_dim = true;
       } else {
+        xla::DExpr size_expr;
         OP_REQUIRES(ctx, size >= 0,
                     errors::InvalidArgument(
                         "size ", d, " must be non-negative, not ", size));
         shape.AddDim(size);
+        if (d < input_shape.dims() && input_shape.get_filled_expression(d) &&
+            input_shape.get_filled_expression(d)->is_dynamic()) {
+          int old = input_shape.dim_size(d);
+          bool is_split = (old > size);
+          int local_ratio = ratio * (is_split ? old / size : size / old);
+	          xla::DExpr input_dexpr = input_shape.get_filled_expression(d);
+	          xla::DExpr ratio_expr = xla::DExpr::Const(local_ratio);
+	          xla::DExpr new_expr =
+	              (size > old) ? input_dexpr * ratio_expr  // Reduce [x,y] -> [x*y]
+	                           : input_dexpr / ratio_expr;  // Split [xy] -> [x/y,y]
+
+          // Pass ratio to next dimension if this is a split, otherwise just
+          // reset it to 1.
+          ratio = is_split ? local_ratio : 1;
+          size_expr = new_expr.simplify();
+
+        } else {
+          size_expr = xla::DExpr::Const(size);
+          if (ratio != 1) {
+            // A split dynamic dimension can be materialized by multiple later
+            // known dimensions. Any unresolved remainder is kept in `ratio`
+            // and may be consumed by a subsequent `-1` dimension (if present);
+            // otherwise, it remains unapplied.
+            if (ratio % size == 0) {
+              ratio /= size;
+            } else if (size % ratio == 0) {
+              ratio = 1;
+            }
+          }
+        }
         product *= size;
+        product_expr = product_expr * size_expr;
+        shape.AddExpression(size_expr);
       }
     }
     auto input = ctx->Input(0);
     if (unknown_index != -1) {
       int64_t input_num_elements = 1;
+      xla::DExpr input_num_elements_expr = xla::DExpr::Const(1);
       bool input_has_zero_dim = false;
       for (int dim = 0; dim < input_shape.dims(); dim++) {
         // For zero dimension, we don't count it into `input_num_elements`
@@ -92,12 +131,18 @@ class ReshapeOp : public XlaOpKernel {
         // infer shapes for other dimensions.
         if (input_shape.dim_size(dim) > 0 || !shape_has_zero_dim) {
           input_num_elements *= input_shape.dim_size(dim);
+          input_num_elements_expr =
+              (input_num_elements_expr * input_shape.get_filled_expression(dim))
+                  .simplify();
         } else {
           input_has_zero_dim = true;
         }
       }
 
       int64_t missing = input_num_elements / product;
+      input_num_elements_expr = input_num_elements_expr.simplify();
+      product_expr = product_expr.simplify();
+      auto missing_expr = input_num_elements_expr / product_expr;
       if (!input_has_zero_dim) {
         if (input_xla_shape->is_static() ||
             input_xla_shape->dimensions().size() != 1) {
@@ -119,10 +164,19 @@ class ReshapeOp : public XlaOpKernel {
               input, xla::Zero(ctx->builder(), input_xla_shape->element_type()),
               0, 0, padded_input_num - input_num_elements);
           input_shape.set_dim(0, padded_input_num);
+          // This expression only approximates the padded size: the true value
+          // uses ceil(input_num_elements / product) * product, which we do not
+          // model symbolically here.
+          xla::DExpr padded_input_num_expr =
+              ((input_num_elements_expr / product_expr) * product_expr)
+                  .simplify();
+          input_shape.set_expression(0, padded_input_num_expr);
         }
       }
       shape.set_dim(unknown_index, missing);
+      shape.set_expression(unknown_index, missing_expr.simplify());
     }
+
     OP_REQUIRES(ctx, shape.num_elements() == input_shape.num_elements(),
                 errors::InvalidArgument("Input to reshape is a tensor with ",
                                         input_shape.num_elements(),
@@ -131,19 +185,23 @@ class ReshapeOp : public XlaOpKernel {
 
     VLOG(2) << "Reshape from " << input_shape.DebugString() << " to "
             << shape.DebugString() << ", unknown_index=" << unknown_index;
+
     if (input_xla_shape->is_static()) {
-      ctx->SetOutput(0, xla::Reshape(input, shape.dim_sizes()));
+      ctx->SetOutput(
+          0, xla::Reshape(input, shape.dim_sizes(), shape.get_filled_expressions()));
       return;
     }
 
     std::vector<xla::XlaOp> output_dim_sizes;
     std::vector<bool> dims_are_dynamic;
+    std::vector<xla::DExpr> output_dim_exprs;
     const auto& dims = shape.dims();
     dims_are_dynamic.reserve(dims);
     output_dim_sizes.reserve(dims);
     for (int64_t i = 0; i < dims; ++i) {
       output_dim_sizes.push_back(
           xla::Reshape(xla::Slice(ctx->Input(1), {i}, {i + 1}, {1}), {}));
+      output_dim_exprs.push_back(xla::DExpr::Unknown(111));
     }
     OP_REQUIRES_OK(
         ctx, ctx->ResolveInputDynamismIntoPredVector(1, &dims_are_dynamic));
@@ -151,7 +209,7 @@ class ReshapeOp : public XlaOpKernel {
       // No unknown index.
       ctx->SetOutput(
           0, xla::DynamicReshape(input, output_dim_sizes, shape.dim_sizes(),
-                                 dims_are_dynamic));
+                                 dims_are_dynamic, output_dim_exprs));
       return;
     }
     auto common_factors =
@@ -166,21 +224,27 @@ class ReshapeOp : public XlaOpKernel {
       // reshape(Tensor([2, 3, 3]), [3, -1, 3]) product of the group
       // containing -1 will be 6.
       xla::XlaOp product = xla::One(ctx->builder(), xla::S32);
+      xla::DExpr product_expr = xla::DExpr::Const(1);
       for (int64_t dim = start.first; dim < end.first; ++dim) {
         if (input_xla_shape->is_dynamic_dimension(dim)) {
           input_is_dynamic = true;
         }
         product = xla::Mul(product, xla::GetDimensionSize(input, dim));
+        product_expr =
+            (product_expr * input_shape.get_filled_expression(dim)).simplify();
       }
       bool unknown_dim_in_group = false;
       // The real size for the -1 dimension in a reshape. E.g., in
       // reshape(Tensor([2, 3, 3]), [3, -1, 3]) this will be 2.
       xla::XlaOp unknown_dim_size = product;
+      xla::DExpr unknown_dim_expr = product_expr;
       for (int64_t dim = start.second; dim < end.second; ++dim) {
         if (dim == unknown_index) {
           unknown_dim_in_group = true;
         } else {
           unknown_dim_size = xla::Div(unknown_dim_size, output_dim_sizes[dim]);
+          unknown_dim_expr =
+              (unknown_dim_expr / output_dim_exprs[dim]).simplify();
         }
       }
 
@@ -188,12 +252,13 @@ class ReshapeOp : public XlaOpKernel {
         // If input dim is dynamic, output dim at the -1 position must be
         // dynamic. Similarly, if input dim is static, output dim has to be
         // static at the -1 dimension.
+        output_dim_exprs[unknown_index] = unknown_dim_expr;
         dims_are_dynamic[unknown_index] = input_is_dynamic;
         output_dim_sizes[unknown_index] = unknown_dim_size;
 
         ctx->SetOutput(
             0, xla::DynamicReshape(input, output_dim_sizes, shape.dim_sizes(),
-                                   dims_are_dynamic));
+                                   dims_are_dynamic, output_dim_exprs));
         VLOG(2) << "Reshape from " << ctx->InputXlaShape(0)->ToString()
                 << " to " << xla::VectorString(shape.dim_sizes())
                 << ", dynamic_dims=" << xla::VectorString(dims_are_dynamic);
diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
index 5cecbf37706283..9f3bbf333eb4eb 100644
--- a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
@@ -71,6 +71,9 @@ class ReverseSequenceOp : public XlaOpKernel {
     xla::XlaBuilder* builder = context->builder();
     const auto input = context->Input(0);
     const auto seq_lens = context->Input(1);
+    auto input_xla_shape_or = context->InputXlaShape(0);
+    OP_REQUIRES(context, input_xla_shape_or.ok(), input_xla_shape_or.status());
+    const xla::Shape& input_xla_shape = input_xla_shape_or.value();
 
     const int64_t batch_size = input_shape.dim_size(batch_dim_);
     if (batch_size == 0) {
@@ -86,11 +89,21 @@ class ReverseSequenceOp : public XlaOpKernel {
     xla::XlaOp back = xla::Sub(seq_lens, xla::ScalarLike(seq_lens, 1));
     xla::XlaOp batch_idx = xla::Iota(
         builder,
-        xla::ShapeUtil::MakeShape(seq_lens_type, {batch_size, max_seq_len, 1}),
+        xla::ShapeUtil::MakeShape(
+            seq_lens_type, {batch_size, max_seq_len, 1},
+            std::vector<xla::DExpr>{
+                input_shape.get_filled_expression(batch_dim_),
+                input_shape.get_filled_expression(seq_dim_),
+                xla::DExpr::Const(1)}),
         /*iota_dimension=*/0);
     xla::XlaOp forward_idx = xla::Iota(
         builder,
-        xla::ShapeUtil::MakeShape(seq_lens_type, {batch_size, max_seq_len, 1}),
+        xla::ShapeUtil::MakeShape(
+            seq_lens_type, {batch_size, max_seq_len, 1},
+            std::vector<xla::DExpr>{
+                input_shape.get_filled_expression(batch_dim_),
+                input_shape.get_filled_expression(seq_dim_),
+                xla::DExpr::Const(1)}),
         /*iota_dimension=*/1);
     xla::XlaOp reverse_idx = xla::Sub(back, forward_idx, {0});
     reverse_idx = xla::Select(xla::Lt(reverse_idx, xla::ZerosLike(reverse_idx)),
@@ -129,8 +142,8 @@ class ReverseSequenceOp : public XlaOpKernel {
     slice_sizes[batch_dim_] = 1;
     slice_sizes[seq_dim_] = 1;
 
-    context->SetOutput(0,
-                       xla::Gather(input, start_indices, dnums, slice_sizes));
+    xla::XlaOp gathered = xla::Gather(input, start_indices, dnums, slice_sizes);
+    context->SetOutput(0, xla::Reshape(input_xla_shape, gathered));
   }
 
  private:
diff --git a/tensorflow/compiler/tf2xla/kernels/roll_op.cc b/tensorflow/compiler/tf2xla/kernels/roll_op.cc
index 0fcc6bec56095b..49b5cd3f01b8b8 100644
--- a/tensorflow/compiler/tf2xla/kernels/roll_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/roll_op.cc
@@ -94,8 +94,8 @@ class RollOp : public XlaOpKernel {
       std::vector<xla::XlaOp> start_indices(
           input_shape.dims(), xla::Zero(ctx->builder(), shift_type));
       start_indices[cur_axis] = axis_size - offset;
-      output =
-          xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes());
+      output = xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes(),
+                                 input_shape.get_filled_expressions());
     }
     ctx->SetOutput(0, output);
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc
index 694b4eb17ef298..f5e03242666a1a 100644
--- a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc
@@ -118,7 +118,8 @@ class ScatterNdOp : public XlaOpKernel {
 
     xla::XlaBuilder* builder = context->builder();
     auto buffer = xla::Broadcast(XlaHelpers::Zero(builder, dtype),
-                                 buffer_shape.dim_sizes());
+                                 buffer_shape.dim_sizes(),
+                                 buffer_shape.get_filled_expressions());
     auto indices = context->Input(0);
     auto updates = context->Input(1);
     auto combine =
diff --git a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc
index 21eaac25f058ed..f731727f582411 100644
--- a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc
@@ -96,7 +96,8 @@ class SegmentReduce : public XlaOpKernel {
     buffer_shape.InsertDim(0, num_segments);
 
     auto buffer =
-        xla::Broadcast(InitialValue(builder), buffer_shape.dim_sizes());
+        xla::Broadcast(InitialValue(builder), buffer_shape.dim_sizes(),
+                       buffer_shape.get_filled_expressions());
 
     // Build dynamic dim sizes for buffer, as well as whether each dimension
     // size is dynamic or static. We build two parts: num_sgement part and
diff --git a/tensorflow/compiler/tf2xla/kernels/select_op.cc b/tensorflow/compiler/tf2xla/kernels/select_op.cc
index 85aaabe87076c2..f6297e9572f76a 100644
--- a/tensorflow/compiler/tf2xla/kernels/select_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/select_op.cc
@@ -70,8 +70,10 @@ class SelectOp : public XlaOpKernel {
       // Broadcast into the dimensions on the right.
       std::vector<int64_t> broadcast_dimensions(cond_shape.dims());
       absl::c_iota(broadcast_dimensions, 0);
+
       cond_handle = xla::BroadcastInDim(cond_handle, then_shape.dim_sizes(),
-                                        broadcast_dimensions);
+                                        broadcast_dimensions,
+                                        then_shape.get_filled_expressions());
     }
     ctx->SetOutput(0, xla::Select(cond_handle, then_handle, else_handle));
   }
@@ -81,7 +83,8 @@ class SelectOp : public XlaOpKernel {
   void operator=(const SelectOp&) = delete;
 };
 
-REGISTER_XLA_OP(Name("Select"), MlirXlaOpKernel);
+// REGISTER_XLA_OP(Name("Select"), MlirXlaOpKernel);
+REGISTER_XLA_OP(Name("Select"), SelectOp);
 
 class SelectOpV2 : public XlaOpKernel {
  public:
diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
index 108bf3848aae93..cd011016a4672c 100644
--- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
@@ -17,8 +17,9 @@ limitations under the License.
 
 #include <cstdint>
 #include <type_traits>
-
+#include "absl/algorithm/container.h"
 #include "absl/status/statusor.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/constants.h"
@@ -26,6 +27,7 @@ limitations under the License.
 #include "xla/hlo/builder/xla_builder.h"
 #include "xla/literal.h"
 #include "xla/primitive_util.h"
+#include "xla/shape_util.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
@@ -33,17 +35,90 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
 namespace {
 
+template <typename T>
+xla::DExpr GetScalarExpr(const XlaExpression& expression,
+                         const xla::LiteralSlice& literal) {
+  const auto& contents = expression.contents();
+  if (!contents.empty() && contents[0]) {
+    return contents[0];
+  }
+  return xla::DExpr::Const(literal.Get<T>({}));
+}
+
+bool HasDynamicContent(const XlaExpression& expression) {
+  return absl::c_any_of(expression.contents(), [](const xla::DExpr& expr) {
+    return expr && expr->is_dynamic();
+  });
+}
+
+template <typename T>
+std::vector<xla::DExpr> BuildRangeContents(const XlaExpression& start_expr,
+                                           const XlaExpression& delta_expr,
+                                           const xla::LiteralSlice& start,
+                                           const xla::LiteralSlice& delta,
+                                           int64_t size) {
+  std::vector<xla::DExpr> contents;
+  contents.reserve(size);
+  xla::DExpr start_symbol = GetScalarExpr<T>(start_expr, start);
+  xla::DExpr delta_symbol = GetScalarExpr<T>(delta_expr, delta);
+  for (int64_t i = 0; i < size; ++i) {
+    xla::DExpr offset = xla::DExpr::Const(static_cast<T>(i));
+    contents.push_back((start_symbol + (delta_symbol * offset).simplify()).simplify());
+  }
+  return contents;
+}
+
+template <typename T>
+xla::DExpr BuildRangeSizeExpr(const XlaExpression& start_expr,
+                              const XlaExpression& limit_expr,
+                              const XlaExpression& delta_expr,
+                              const xla::LiteralSlice& start,
+                              const xla::LiteralSlice& limit,
+                              const xla::LiteralSlice& delta,
+                              int64_t fallback_size) {
+  xla::DExpr start_symbol = GetScalarExpr<T>(start_expr, start);
+  xla::DExpr limit_symbol = GetScalarExpr<T>(limit_expr, limit);
+  xla::DExpr delta_symbol = GetScalarExpr<T>(delta_expr, delta);
+
+  const auto& start_contents = start_expr.contents();
+  xla::DExpr effective_start =
+      (!start_contents.empty() && start_contents[0]) ? start_contents[0]
+                                                     : start_symbol;
+  const auto& limit_contents = limit_expr.contents();
+  xla::DExpr effective_limit =
+      (!limit_contents.empty() && limit_contents[0]) ? limit_contents[0]
+                                                     : limit_symbol;
+  const auto& delta_contents = delta_expr.contents();
+  xla::DExpr effective_delta =
+      (!delta_contents.empty() && delta_contents[0]) ? delta_contents[0]
+                                                     : delta_symbol;
+
+  xla::DExpr positive_diff = (effective_limit - effective_start).simplify();
+  xla::DExpr positive_size =
+      (((positive_diff - 1) / effective_delta) + 1).simplify();
+  xla::DExpr negative_step =
+      (xla::DExpr::Const(0) - effective_delta).simplify();
+  xla::DExpr negative_diff = (effective_start - effective_limit).simplify();
+  xla::DExpr negative_size =
+      (((negative_diff - 1) / negative_step) + 1).simplify();
+  return xla::DExpr::Select(xla::DExpr::Gt(delta_symbol, xla::DExpr::Const(0)),
+                            positive_size, negative_size)
+      .simplify();
+}
+
 // The type-specific part of the implementation of Range.
 template <typename T>
 absl::StatusOr<xla::XlaOp> CreateRangeTensor(
     const xla::LiteralSlice& start_literal,
     const xla::LiteralSlice& limit_literal,
-    const xla::LiteralSlice& delta_literal, xla::XlaBuilder* builder) {
+    const xla::LiteralSlice& delta_literal, xla::XlaBuilder* builder,
+    xla::DExpr size_expr = xla::DExpr()) {
   T start = start_literal.Get<T>({});
   T limit = limit_literal.Get<T>({});
   T delta = delta_literal.Get<T>({});
@@ -70,10 +145,17 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
                      : (std::abs(limit - start) - 1) / std::abs(delta) + 1)
            : std::ceil(std::abs((limit - start) / delta)));
 
-  return xla::ConstantR0(builder, start) +
-         xla::ConstantR0(builder, delta) *
-             xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType<T>(),
-                       size);
+  xla::XlaOp iota =
+      (std::is_integral<T>::value && size_expr)
+          ? xla::Iota(builder,
+                      xla::ShapeUtil::MakeShape(
+                          xla::primitive_util::NativeToPrimitiveType<T>(),
+                          {size}, std::vector<xla::DExpr>{size_expr}),
+                      /*iota_dimension=*/0)
+          : xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType<T>(),
+                      size);
+
+  return xla::ConstantR0(builder, start) + xla::ConstantR0(builder, delta) * iota;
 }
 
 class RangeOp : public XlaOpKernel {
@@ -103,13 +185,50 @@ class RangeOp : public XlaOpKernel {
     DataType type = input_type(0);
     absl::StatusOr<xla::XlaOp> output;
     switch (type) {
-      case DT_INT32:
-        output = CreateRangeTensor<int32>(start, limit, delta, ctx->builder());
+      case DT_INT32: {
+        int32 start_value = start.Get<int32>({});
+        int32 limit_value = limit.Get<int32>({});
+        int32 delta_value = delta.Get<int32>({});
+        int64_t size = static_cast<int32>(
+            limit_value == start_value
+                ? 0
+                : (std::abs(limit_value - start_value) - 1) /
+                          std::abs(delta_value) +
+                      1);
+        xla::DExpr size_expr = xla::DExpr::Const(size);
+        if (HasDynamicContent(ctx->InputExpression(0)) ||
+            HasDynamicContent(ctx->InputExpression(1)) ||
+            HasDynamicContent(ctx->InputExpression(2))) {
+          size_expr = BuildRangeSizeExpr<int32>(
+              ctx->InputExpression(0), ctx->InputExpression(1),
+              ctx->InputExpression(2), start, limit, delta, size);
+        }
+        output = CreateRangeTensor<int32>(start, limit, delta, ctx->builder(),
+                                          size_expr);
         break;
-      case DT_INT64:
-        output =
-            CreateRangeTensor<int64_t>(start, limit, delta, ctx->builder());
+      }
+      case DT_INT64: {
+        int64_t start_value = start.Get<int64_t>({});
+        int64_t limit_value = limit.Get<int64_t>({});
+        int64_t delta_value = delta.Get<int64_t>({});
+        int64_t size =
+            limit_value == start_value
+                ? 0
+                : (std::abs(limit_value - start_value) - 1) /
+                          std::abs(delta_value) +
+                      1;
+        xla::DExpr size_expr = xla::DExpr::Const(size);
+        if (HasDynamicContent(ctx->InputExpression(0)) ||
+            HasDynamicContent(ctx->InputExpression(1)) ||
+            HasDynamicContent(ctx->InputExpression(2))) {
+          size_expr = BuildRangeSizeExpr<int64_t>(
+              ctx->InputExpression(0), ctx->InputExpression(1),
+              ctx->InputExpression(2), start, limit, delta, size);
+        }
+        output = CreateRangeTensor<int64_t>(start, limit, delta, ctx->builder(),
+                                            size_expr);
         break;
+      }
       case DT_FLOAT:
         output = CreateRangeTensor<float>(start, limit, delta, ctx->builder());
         break;
@@ -145,7 +264,55 @@ class RangeOp : public XlaOpKernel {
       }
     }
 
-    ctx->SetOutput(0, output.value());
+    const XlaExpression& start_expr = ctx->InputExpression(0);
+    const XlaExpression& limit_expr = ctx->InputExpression(1);
+    const XlaExpression& delta_expr = ctx->InputExpression(2);
+    const bool symbolic_enabled = SymbolicContentEnabled();
+    const bool has_dynamic_content =
+        HasDynamicContent(start_expr) || HasDynamicContent(limit_expr) ||
+        HasDynamicContent(delta_expr);
+
+    if (type == DT_INT32) {
+      int32 start_value = start.Get<int32>({});
+      int32 limit_value = limit.Get<int32>({});
+      int32 delta_value = delta.Get<int32>({});
+      if (symbolic_enabled && has_dynamic_content) {
+        int64_t size = static_cast<int32>(
+            limit_value == start_value
+                ? 0
+                : (std::abs(limit_value - start_value) - 1) /
+                          std::abs(delta_value) +
+                      1);
+        auto output_expr =
+            XlaExpression::XlaOp(output.value(), ctx->expected_output_dtype(0));
+        output_expr.set_contents(BuildRangeContents<int32>(
+            start_expr, delta_expr, start, delta, size));
+        ctx->SetOutputExpression(0, output_expr);
+      } else {
+        ctx->SetOutput(0, output.value());
+      }
+    } else if (type == DT_INT64) {
+      int64_t start_value = start.Get<int64_t>({});
+      int64_t limit_value = limit.Get<int64_t>({});
+      int64_t delta_value = delta.Get<int64_t>({});
+      if (symbolic_enabled && has_dynamic_content) {
+        int64_t size = static_cast<int64_t>(
+            limit_value == start_value
+                ? 0
+                : (std::abs(limit_value - start_value) - 1) /
+                          std::abs(delta_value) +
+                      1);
+        auto output_expr =
+            XlaExpression::XlaOp(output.value(), ctx->expected_output_dtype(0));
+        output_expr.set_contents(BuildRangeContents<int64_t>(
+            start_expr, delta_expr, start, delta, size));
+        ctx->SetOutputExpression(0, output_expr);
+      } else {
+        ctx->SetOutput(0, output.value());
+      }
+    } else {
+      ctx->SetOutput(0, output.value());
+    }
   }
 };
 
diff --git a/tensorflow/compiler/tf2xla/kernels/shape_op.cc b/tensorflow/compiler/tf2xla/kernels/shape_op.cc
index 7e8889cb2ccee6..1d624598d6ba07 100644
--- a/tensorflow/compiler/tf2xla/kernels/shape_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/shape_op.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "tensorflow/compiler/tf2xla/kernels/shape_util.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/kernels/tensor_list_utils.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
@@ -46,6 +47,19 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+std::vector<xla::DExpr> BuildShapeContents(const TensorShape& input_shape) {
+  std::vector<xla::DExpr> contents;
+  contents.reserve(input_shape.dims());
+  for (int64_t i = 0; i < input_shape.dims(); ++i) {
+    xla::DExpr expr = input_shape.get_filled_expression(i);
+    contents.push_back(expr && expr->is_dynamic()
+                           ? expr
+                           : xla::DExpr::Unknown(
+                                 xla::kUnknownContentSentinel));
+  }
+  return contents;
+}
+
 class ShapeOp : public XlaOpKernel {
  public:
   explicit ShapeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
@@ -58,18 +72,46 @@ class ShapeOp : public XlaOpKernel {
     const int rank = input_shape.dims();
     if (rank != 0) {
       for (int64_t i = 0; i < rank; ++i) {
-        operands.push_back(xla::Broadcast(
-            xla::ConvertElementType(xla::GetDimensionSize(ctx->Input(0), i),
-                                    ctx->output_xla_type(0)),
-            {1}));
+        xla::DExpr expr = input_shape.get_filled_expression(i);
+        std::vector<xla::DExpr> content = {
+            expr && expr->is_dynamic()
+                ? expr
+                : xla::DExpr::Unknown(xla::kUnknownContentSentinel)};
+        xla::XlaOp dim_size = xla::GetDimensionSize(ctx->Input(0), i);
+        if (SymbolicContentEnabled()) {
+          OP_REQUIRES_OK(
+              ctx, ctx->builder()->SetInstructionContents(dim_size, content));
+        }
+        xla::XlaOp converted =
+            xla::ConvertElementType(dim_size, ctx->output_xla_type(0));
+        if (SymbolicContentEnabled()) {
+          OP_REQUIRES_OK(
+              ctx, ctx->builder()->SetInstructionContents(converted, content));
+        }
+        xla::XlaOp broadcast = xla::Broadcast(converted, {1});
+        if (SymbolicContentEnabled()) {
+          OP_REQUIRES_OK(
+              ctx, ctx->builder()->SetInstructionContents(broadcast, content));
+        }
+        operands.push_back(broadcast);
       }
 
-      ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), operands, 0));
+      xla::XlaOp concat = xla::ConcatInDim(ctx->builder(), operands, 0);
+      XlaExpression output =
+          XlaExpression::XlaOp(concat, ctx->expected_output_dtype(0));
+      if (SymbolicContentEnabled()) {
+        output.set_contents(BuildShapeContents(input_shape));
+      }
+      ctx->SetOutputExpression(0, output);
     } else {
       // Rank 0 won't have dynamic size dimension, use constant output.
       Tensor shape_constant(out_dtype_, TensorShape({input_shape.dims()}));
       OP_REQUIRES_OK(ctx, TensorShapeToConstant(input_shape, &shape_constant));
-      ctx->SetConstantOutput(0, shape_constant);
+      XlaExpression output = XlaExpression::Constant(shape_constant);
+      if (SymbolicContentEnabled()) {
+        output.set_contents(BuildShapeContents(input_shape));
+      }
+      ctx->SetOutputExpression(0, output);
     }
   }
 
@@ -196,19 +238,47 @@ class ShapeNOp : public XlaOpKernel {
         // Each dimension can be dynamic, so use GetDimensionSize to get the
         // runtime dimension.
         for (int64_t dim = 0; dim < rank; ++dim) {
-          operands.push_back(xla::Broadcast(
-              xla::ConvertElementType(xla::GetDimensionSize(ctx->Input(i), dim),
-                                      ctx->output_xla_type(i)),
-              {1}));
+          xla::DExpr expr = input_shape.get_filled_expression(dim);
+          std::vector<xla::DExpr> content = {
+              expr && expr->is_dynamic()
+                  ? expr
+                  : xla::DExpr::Unknown(xla::kUnknownContentSentinel)};
+          xla::XlaOp dim_size = xla::GetDimensionSize(ctx->Input(i), dim);
+          if (SymbolicContentEnabled()) {
+            OP_REQUIRES_OK(
+                ctx, ctx->builder()->SetInstructionContents(dim_size, content));
+          }
+          xla::XlaOp converted =
+              xla::ConvertElementType(dim_size, ctx->output_xla_type(i));
+          if (SymbolicContentEnabled()) {
+            OP_REQUIRES_OK(
+                ctx, ctx->builder()->SetInstructionContents(converted, content));
+          }
+          xla::XlaOp broadcast = xla::Broadcast(converted, {1});
+          if (SymbolicContentEnabled()) {
+            OP_REQUIRES_OK(
+                ctx, ctx->builder()->SetInstructionContents(broadcast, content));
+          }
+          operands.push_back(broadcast);
         }
 
-        ctx->SetOutput(i, xla::ConcatInDim(ctx->builder(), operands, 0));
+        XlaExpression output =
+            XlaExpression::XlaOp(xla::ConcatInDim(ctx->builder(), operands, 0),
+                                 ctx->expected_output_dtype(i));
+        if (SymbolicContentEnabled()) {
+          output.set_contents(BuildShapeContents(input_shape));
+        }
+        ctx->SetOutputExpression(i, output);
       } else {
         // Rank 0 won't have dynamic size dimension, use constant output.
         Tensor shape_constant(out_dtype_, TensorShape({input_shape.dims()}));
         OP_REQUIRES_OK(ctx,
                        TensorShapeToConstant(input_shape, &shape_constant));
-        ctx->SetConstantOutput(i, shape_constant);
+        XlaExpression output = XlaExpression::Constant(shape_constant);
+        if (SymbolicContentEnabled()) {
+          output.set_contents(BuildShapeContents(input_shape));
+        }
+        ctx->SetOutputExpression(i, output);
       }
     }
   }
@@ -244,6 +314,7 @@ class SizeOp : public XlaOpKernel {
     const TensorShape input_shape = ctx->InputShape(0);
     xla::XlaBuilder* builder = ctx->builder();
     auto size = xla::One(builder, ctx->output_xla_type(0));
+    xla::DExpr size_expr = xla::DExpr::Const(1);
 
     const int rank = input_shape.dims();
     for (int64_t dim = 0; dim < rank; ++dim) {
@@ -256,12 +327,39 @@ class SizeOp : public XlaOpKernel {
               "on all dimensions, found ",
               input_shape.dim_size(dim), " elements on dimension ", dim)));
 
-      size = xla::Mul(size, xla::ConvertElementType(
-                                xla::GetDimensionSize(ctx->Input(0), dim),
-                                ctx->output_xla_type(0)));
+      xla::DExpr dim_expr = input_shape.get_filled_expression(dim);
+      std::vector<xla::DExpr> contents = {
+          dim_expr && dim_expr->is_dynamic()
+              ? dim_expr
+              : xla::DExpr::Unknown(xla::kUnknownContentSentinel)};
+      xla::XlaOp dim_size = xla::GetDimensionSize(ctx->Input(0), dim);
+      if (SymbolicContentEnabled()) {
+        OP_REQUIRES_OK(ctx,
+                       builder->SetInstructionContents(dim_size, contents));
+      }
+      xla::XlaOp converted =
+          xla::ConvertElementType(dim_size, ctx->output_xla_type(0));
+      if (SymbolicContentEnabled()) {
+        OP_REQUIRES_OK(ctx,
+                       builder->SetInstructionContents(converted, contents));
+      }
+      size = xla::Mul(size, converted);
+      size_expr =
+          (size_expr * input_shape.get_filled_expression(dim)).simplify();
+      if (SymbolicContentEnabled()) {
+        OP_REQUIRES_OK(ctx,
+                       builder->SetInstructionContents(size, {size_expr}));
+      }
     }
 
-    ctx->SetOutput(0, size);
+    if (SymbolicContentEnabled()) {
+      XlaExpression output =
+          XlaExpression::XlaOp(size, ctx->expected_output_dtype(0));
+      output.set_contents({size_expr});
+      ctx->SetOutputExpression(0, output);
+    } else {
+      ctx->SetOutput(0, size);
+    }
   }
 };
 
@@ -290,6 +388,8 @@ class ExpandDimsOp : public XlaOpKernel {
                                         " dimensions."));
 
     auto existing_dims = input_shape.dim_sizes();
+    auto existing_exprs = input_shape.get_filled_expressions();
+
     // Safe - # elements in tensor dims bounded.
     const int existing_dims_size = static_cast<int>(existing_dims.size());
     std::vector<int64_t> new_shape(existing_dims_size);
@@ -297,6 +397,13 @@ class ExpandDimsOp : public XlaOpKernel {
       new_shape[i] = existing_dims[i];
     }
 
+    const int existing_exprs_size = static_cast<int>(existing_exprs.size());
+    std::vector<xla::DExpr> new_exprs;
+    new_exprs.reserve(existing_exprs_size);
+    for (size_t i = 0; i < existing_exprs.size(); ++i) {
+      new_exprs.push_back(existing_exprs[i]);
+    }
+
     // We emulate numpy's interpretation of the dim axis when
     // -input.dims() >= dim <= input.dims().
     if (dim < 0) {
@@ -306,8 +413,8 @@ class ExpandDimsOp : public XlaOpKernel {
     // Clamp to the end if needed.
     dim = std::min<int32_t>(dim, existing_dims_size);
     new_shape.emplace(new_shape.begin() + dim, 1);
-
-    ctx->SetOutput(0, xla::Reshape(ctx->Input("input"), new_shape));
+    new_exprs.emplace(new_exprs.begin() + dim, xla::DExpr::Const(1));
+    ctx->SetOutput(0, xla::Reshape(ctx->Input("input"), new_shape, new_exprs));
   }
 };
 REGISTER_XLA_OP(Name("ExpandDims").CompileTimeConstantInput("dim"),
@@ -331,6 +438,7 @@ class SqueezeOp : public XlaOpKernel {
     absl::flat_hash_set<int32_t> wrapped_squeeze_dims;
     wrapped_squeeze_dims.reserve(squeeze_dims_.size());
     std::vector<int64_t> new_shape;
+    std::vector<xla::DExpr> new_exprs;
     // Validate squeeze dims against the input.
     for (int32_t dim : squeeze_dims_) {
       OP_REQUIRES(
@@ -358,6 +466,7 @@ class SqueezeOp : public XlaOpKernel {
         } else {
           // This dimension is not being squeezed.
           new_shape.push_back(existing_dim);
+          new_exprs.push_back(shape.expressions(i));
         }
       } else {
         OP_REQUIRES(
@@ -368,11 +477,25 @@ class SqueezeOp : public XlaOpKernel {
         // Copy over all non-1-length dimensions.
         if (existing_dim != 1) {
           new_shape.push_back(existing_dim);
+          new_exprs.push_back(shape.expressions(i));
         }
       }
     }
 
-    ctx->SetOutput(0, xla::Reshape(ctx->Input(0), new_shape));
+    xla::XlaOp output = xla::Reshape(ctx->Input(0), new_shape, new_exprs);
+    const auto& input_contents = ctx->InputExpression(0).contents();
+    if (SymbolicContentEnabled() && !input_contents.empty()) {
+      std::vector<xla::DExpr> output_contents(input_contents.begin(),
+                                              input_contents.end());
+      OP_REQUIRES_OK(
+          ctx, ctx->builder()->SetInstructionContents(output, output_contents));
+      XlaExpression output_expr =
+          XlaExpression::XlaOp(output, ctx->expected_output_dtype(0));
+      output_expr.set_contents(std::move(output_contents));
+      ctx->SetOutputExpression(0, output_expr);
+      return;
+    }
+    ctx->SetOutput(0, output);
   }
 
  private:
@@ -430,7 +553,8 @@ class ZerosLikeOp : public XlaOpKernel {
       auto zero = XlaHelpers::Zero(ctx->builder(), input_type(0));
       xla::XlaOp input = ctx->Input(0);
       auto input_shape = ctx->InputXlaShape(0).value();
-      auto result = xla::Broadcast(zero, input_shape.dimensions());
+      auto result = xla::Broadcast(zero, input_shape.dimensions(),
+                                   input_shape.expressions());
 
       // Setting up dynamic dimensions of the broadcast.
       for (int64_t i = 0; i < input_shape.dimensions().size(); ++i) {
@@ -455,7 +579,8 @@ class OnesLikeOp : public XlaOpKernel {
     const TensorShape input_shape = ctx->InputShape(0);
 
     auto one = XlaHelpers::One(ctx->builder(), input_type(0));
-    ctx->SetOutput(0, xla::Broadcast(one, input_shape.dim_sizes()));
+    ctx->SetOutput(0, xla::Broadcast(one, input_shape.dim_sizes(),
+                                     input_shape.get_filled_expressions()));
   }
 };
 
diff --git a/tensorflow/compiler/tf2xla/kernels/slice_op.cc b/tensorflow/compiler/tf2xla/kernels/slice_op.cc
index 844a31f97990fc..5d3d9023bf5a02 100644
--- a/tensorflow/compiler/tf2xla/kernels/slice_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/slice_op.cc
@@ -18,8 +18,11 @@ limitations under the License.
 #include <cstdint>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/types/span.h"
+#include "tensorflow/compiler/jit/flags.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/constants.h"
@@ -37,11 +40,42 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+bool TryBuildSlicedContents(const XlaExpression& input_expr,
+                            const TensorShape& input_shape,
+                            absl::Span<const int64_t> begin,
+                            absl::Span<const int64_t> size,
+                            std::vector<xla::DExpr>* output_contents) {
+  output_contents->clear();
+  const auto& input_contents = input_expr.contents();
+  if (input_contents.empty() || input_shape.dims() != 1 || begin.size() != 1 ||
+      size.size() != 1) {
+    return false;
+  }
+  const int64_t start = begin[0];
+  const int64_t count =
+      size[0] == -1 ? input_shape.dim_size(0) - start : size[0];
+  for (int64_t i = 0; i < count; ++i) {
+    const int64_t index = start + i;
+    if (index < 0 || index >= input_contents.size()) {
+      output_contents->clear();
+      return false;
+    }
+    const xla::DExpr& expr = input_contents[index];
+    output_contents->push_back(
+        expr ? expr : xla::DExpr::Unknown(xla::kUnknownContentSentinel));
+  }
+  return absl::c_any_of(*output_contents, [](const xla::DExpr& expr) {
+    return expr && expr->is_dynamic();
+  });
+}
+
 class SliceOp : public XlaOpKernel {
  public:
   explicit SliceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
+    const bool enable_dynamic_sizes =
+        GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes;
     const TensorShape input_shape = ctx->InputShape(0);
     const TensorShape begin_tensor_shape = ctx->InputShape(1);
     const TensorShape size_tensor_shape = ctx->InputShape(2);
@@ -66,13 +100,17 @@ class SliceOp : public XlaOpKernel {
         ctx->ConstantInputAsIntVector(2, &size).ok();
     if (all_begins_are_constant && all_sizes_are_constant) {
       std::vector<int64_t> wrapped_size(size.size());
+      std::vector<xla::DExpr> wrapped_size_exprs(size.size());
       // `begin` is a compile-time constant.
       for (int i = 0; i < input_dims; ++i) {
         if (size[i] == -1) {
           // A size[i] of -1 means "all elements from begin[i] to dim_size(i)".
           wrapped_size[i] = input_shape.dim_size(i) - begin[i];
+          wrapped_size_exprs[i] =
+              (input_shape.get_filled_expression(i) - begin[i]).simplify();
         } else {
           wrapped_size[i] = size[i];
+          wrapped_size_exprs[i] = xla::DExpr::Const(size[i]);
         }
       }
 
@@ -97,13 +135,32 @@ class SliceOp : public XlaOpKernel {
         }
       }
 
+      std::vector<xla::DExpr> begin_exprs;
+      for (int d : begin){
+        begin_exprs.push_back(xla::DExpr::Const(d));
+      }
       std::vector<int64_t> limits;
+      std::vector<xla::DExpr> exprs;
       limits.reserve(begin.size());
+      exprs.reserve(begin.size());
       for (int i = 0; i < begin.size(); ++i) {
         limits.push_back(begin[i] + wrapped_size[i]);
+        exprs.push_back(begin_exprs[i] + wrapped_size_exprs[i]);
       }
       std::vector<int64_t> strides(begin.size(), 1);
-      auto slice = xla::Slice(ctx->Input(0), begin, limits, strides);
+      auto slice = enable_dynamic_sizes
+                       ? xla::Slice(ctx->Input(0), begin, limits, begin_exprs,
+                                    exprs, strides)
+                       : xla::Slice(ctx->Input(0), begin, limits, strides);
+      std::vector<xla::DExpr> output_contents;
+      const bool has_output_contents =
+          SymbolicContentEnabled() &&
+          TryBuildSlicedContents(ctx->InputExpression(0), input_shape, begin,
+                                 size, &output_contents);
+      if (has_output_contents) {
+        OP_REQUIRES_OK(
+            ctx, ctx->builder()->SetInstructionContents(slice, output_contents));
+      }
       // Check for slice on dynamic dimensions.
       std::vector<bool> size_is_dynamic;
       OP_REQUIRES_OK(
@@ -114,13 +171,27 @@ class SliceOp : public XlaOpKernel {
           if (size[i] != -1) {
             // If there is a dynamic dimension, properly set dimension size of
             // the slice.
-            auto dynamic_size =
-                xla::Reshape(xla::Slice(ctx->Input(2), {i}, {i + 1}, {1}), {});
+            auto dynamic_size = xla::Reshape(
+                xla::Slice(ctx->Input(2), {i}, {i + 1}, {xla::DExpr::Const(i)},
+                           {xla::DExpr::Const(i + 1)}, {1}),
+                {});
 
             slice = xla::SetDimensionSize(slice, dynamic_size, i);
+            if (has_output_contents) {
+              OP_REQUIRES_OK(
+                  ctx,
+                  ctx->builder()->SetInstructionContents(slice, output_contents));
+            }
           }
         }
       }
+      if (has_output_contents) {
+        auto output_expr =
+            XlaExpression::XlaOp(slice, ctx->expected_output_dtype(0));
+        output_expr.set_contents(std::move(output_contents));
+        ctx->SetOutputExpression(0, output_expr);
+        return;
+      }
       ctx->SetOutput(0, slice);
     } else {
       // When a size is -1, we take rest of the dimension according to
@@ -153,7 +224,14 @@ class SliceOp : public XlaOpKernel {
       }
       if (all_sizes_are_constant && !constant_size_is_minus_one) {
         xla::XlaOp input = ctx->Input(0);
-        ctx->SetOutput(0, xla::DynamicSlice(input, begin_indices, size));
+        std::vector<xla::DExpr> output_exprs;
+        output_exprs.reserve(size.size());
+        for (int64_t d : size) {
+          output_exprs.push_back(xla::DExpr::Const(d));
+        }
+        ctx->SetOutput(0,
+                       xla::DynamicSlice(input, begin_indices, size,
+                                         output_exprs));
       } else {
         // Size is not constant, use input size as upperbound and then set
         // dimension size on it.
diff --git a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc
index 330479bc8d4150..b1d7e16bbda741 100644
--- a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc
@@ -35,11 +35,72 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/errors.h"
 
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/bcast.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+
 namespace tensorflow {
 namespace {
 
-REGISTER_XLA_OP(Name("Softmax"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("LogSoftmax"), MlirXlaOpKernel);
+class SoftmaxOp : public XlaOpKernel {
+ public:
+  explicit SoftmaxOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
+    log_ = absl::StartsWith(type_string(), "Log");
+  }
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    const TensorShape logits_shape = ctx->InputShape(0);
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVectorOrHigher(logits_shape),
+                errors::InvalidArgument("logits must have >= 1 dimension, got ",
+                                        logits_shape.DebugString()));
+
+    // Major dimensions are batch dimensions, minor dimension is the class
+    // dimension.
+    std::vector<int64> batch_dims(logits_shape.dims() - 1);
+    std::iota(batch_dims.begin(), batch_dims.end(), 0);
+    const int kClassDim = logits_shape.dims() - 1;
+
+    const DataType type = input_type(0);
+    const xla::PrimitiveType xla_type = ctx->input_xla_type(0);
+    auto logits = ctx->Input(0);
+
+    xla::XlaBuilder* const b = ctx->builder();
+
+    const xla::XlaComputation& max_func = *ctx->GetOrCreateMax(type);
+
+    // Find the max in each batch, resulting in a tensor of shape [batch]
+    auto logits_max =
+        xla::Reduce(logits, xla::MinValue(b, xla_type), max_func, {kClassDim});
+    // Subtract the max in batch b from every element in batch b. Broadcasts
+    // along the batch dimension.
+    auto shifted_logits = xla::Sub(logits, logits_max, batch_dims);
+    auto exp_shifted = xla::Exp(shifted_logits);
+    const DataType accumulation_type = XlaHelpers::SumAccumulationType(type);
+    xla::PrimitiveType xla_accumulation_type;
+    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(accumulation_type,
+                                                &xla_accumulation_type));
+    auto converted =
+        xla::ConvertElementType(exp_shifted, xla_accumulation_type);
+    auto reduce =
+        xla::Reduce(converted, xla::Zero(b, xla_accumulation_type),
+                    *ctx->GetOrCreateAdd(accumulation_type), {kClassDim});
+    auto sum = XlaHelpers::ConvertElementType(reduce, type);
+    auto softmax =
+        log_
+            // softmax = shifted_logits - log(sum(exp(shifted_logits)))
+            ? xla::Sub(shifted_logits, xla::Log(sum), batch_dims)
+            // softmax = exp(shifted_logits) / sum(exp(shifted_logits))
+            : xla::Div(exp_shifted, sum, batch_dims);
+    ctx->SetOutput(0, softmax);
+  }
+
+ private:
+  bool log_;
+};
+
+REGISTER_XLA_OP(Name("Softmax"), SoftmaxOp);
+REGISTER_XLA_OP(Name("LogSoftmax"), SoftmaxOp);
 
 std::pair<xla::XlaOp, xla::XlaOp> CrossEntropyWithLogits(
     XlaOpKernelContext* ctx, DataType type, xla::PrimitiveType xla_type,
diff --git a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
index d4a93e0556143d..e10c009a5b7d16 100644
--- a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc
@@ -44,6 +44,8 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input,
   const int input_rank = input_tensor_shape.dims();
   const absl::InlinedVector<int64_t, 4> input_shape =
       input_tensor_shape.dim_sizes();
+  const std::vector<xla::DExpr> input_exprs =
+      input_tensor_shape.get_filled_expressions();
   const int block_rank = block_shape.size();
 
   OP_REQUIRES(
@@ -68,6 +70,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input,
   //  input according to `paddings` to produce `padded` of shape `padded_shape`.
   xla::PaddingConfig padding_config;
   std::vector<int64_t> padded_shape(input_shape.begin(), input_shape.end());
+  std::vector<xla::DExpr> padded_exprs(input_exprs.begin(), input_exprs.end());
   int64_t block_num_elems = 1LL;
   padding_config.add_dimensions();  // Don't pad the batch dimension.
   for (int i = 0; i < block_rank; ++i) {
@@ -83,6 +86,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input,
     dim->set_edge_padding_low(pad_start);
     dim->set_edge_padding_high(pad_end);
     padded_shape[1 + i] += pad_start + pad_end;
+    padded_exprs[1 + i] = (padded_exprs[1 + i] + pad_start + pad_end).simplify();
     block_num_elems = MultiplyWithoutOverflow(block_num_elems, block_shape[i]);
   }
   // Don't pad the remainder dimensions.
@@ -116,7 +120,9 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input,
   //       block_shape[M-1]] +
   //      remaining_shape
   std::vector<int64_t> reshaped_padded_shape(input_rank + block_rank);
+  std::vector<xla::DExpr> reshaped_padded_exprs(input_rank + block_rank);
   reshaped_padded_shape[0] = batch_size;
+  reshaped_padded_exprs[0] = padded_exprs[0];
   for (int i = 0; i < block_rank; ++i) {
     OP_REQUIRES(ctx, padded_shape[1 + i] % block_shape[i] == 0,
                 errors::InvalidArgument("padded_shape[", 1 + i,
@@ -126,11 +132,17 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input,
 
     reshaped_padded_shape[1 + i * 2] = padded_shape[1 + i] / block_shape[i];
     reshaped_padded_shape[1 + i * 2 + 1] = block_shape[i];
+    reshaped_padded_exprs[1 + i * 2] =
+        (padded_exprs[1 + i] / block_shape[i]).simplify();
+    reshaped_padded_exprs[1 + i * 2 + 1] = xla::DExpr::Const(block_shape[i]);
   }
   std::copy(remainder_shape.begin(), remainder_shape.end(),
             reshaped_padded_shape.begin() + 1 + 2 * block_rank);
+  std::copy(input_exprs.begin() + 1 + block_rank, input_exprs.end(),
+            reshaped_padded_exprs.begin() + 1 + 2 * block_rank);
 
-  xla::XlaOp reshaped_padded = xla::Reshape(padded, reshaped_padded_shape);
+  xla::XlaOp reshaped_padded =
+      xla::Reshape(padded, reshaped_padded_shape, reshaped_padded_exprs);
 
   // 3. Permute dimensions of `reshaped_padded` to produce
   //    `permuted_reshaped_padded` of shape:
@@ -163,14 +175,21 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp input,
   // Determine the length of the prefix of block dims that can be combined
   // into the batch dimension due to having no padding and block_shape=1.
   std::vector<int64_t> output_shape(input_rank);
+  std::vector<xla::DExpr> output_exprs(input_rank);
   output_shape[0] = output_dim;
+  output_exprs[0] = (input_exprs[0] * xla::DExpr::Const(block_num_elems)).simplify();
   for (int i = 0; i < block_rank; ++i) {
     output_shape[1 + i] = padded_shape[1 + i] / block_shape[i];
+    output_exprs[1 + i] =
+        (padded_exprs[1 + i] / block_shape[i]).simplify();
   }
   std::copy(remainder_shape.begin(), remainder_shape.end(),
             output_shape.begin() + 1 + block_rank);
+  std::copy(input_exprs.begin() + 1 + block_rank, input_exprs.end(),
+            output_exprs.begin() + 1 + block_rank);
 
-  xla::XlaOp output = xla::Reshape(permuted_reshaped_padded, output_shape);
+  xla::XlaOp output =
+      xla::Reshape(permuted_reshaped_padded, output_shape, output_exprs);
   ctx->SetOutput(0, output);
 }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc
index ac33e0877200dc..d09fa5f4daacbb 100644
--- a/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/spacetodepth_op.cc
@@ -67,6 +67,7 @@ class SpaceToDepthOp : public XlaOpKernel {
     OP_REQUIRES_OK(ctx, input_xla_shape.status());
     absl::Span<const int64_t> input_shape =
         input_xla_shape.value().dimensions();
+    const xla::Shape& input_shape_with_exprs = input_xla_shape.value();
     int input_rank = input_shape.size();
 
     static const int kRequiredDims = 4;
@@ -80,9 +81,13 @@ class SpaceToDepthOp : public XlaOpKernel {
     std::vector<int64_t> reshaped_shape;
     std::vector<int64_t> transpose_order;
     std::vector<int64_t> output_shape;
+    std::vector<xla::DExpr> reshaped_exprs;
+    std::vector<xla::DExpr> output_exprs;
     reshaped_shape.reserve(input_rank);
     transpose_order.reserve(input_rank);
     output_shape.reserve(input_rank);
+    reshaped_exprs.reserve(input_rank + num_spatial_dims);
+    output_exprs.reserve(input_rank);
     if (data_format == FORMAT_NHWC) {
       int64_t block_elems = 1;
       for (int i = 0; i < num_spatial_dims; ++i) {
@@ -94,11 +99,18 @@ class SpaceToDepthOp : public XlaOpKernel {
       }
 
       reshaped_shape.push_back(input_shape[0]);
+      reshaped_exprs.push_back(input_shape_with_exprs.expressions(0));
       for (int i = 0; i < num_spatial_dims; ++i) {
         reshaped_shape.push_back(input_shape[1 + i] / block_size_);
+        reshaped_exprs.push_back(
+            (input_shape_with_exprs.expressions(1 + i) /
+             xla::DExpr::Const(block_size_))
+                .simplify());
         reshaped_shape.push_back(block_size_);
+        reshaped_exprs.push_back(xla::DExpr::Const(block_size_));
       }
       reshaped_shape.push_back(input_shape[feature_dim]);
+      reshaped_exprs.push_back(input_shape_with_exprs.expressions(feature_dim));
 
       transpose_order.push_back(0);
       for (int i = 0; i < num_spatial_dims; ++i) {
@@ -110,10 +122,19 @@ class SpaceToDepthOp : public XlaOpKernel {
       transpose_order.push_back(feature_dim + num_spatial_dims);
 
       output_shape.push_back(input_shape[0]);
+      output_exprs.push_back(input_shape_with_exprs.expressions(0));
       for (int i = 0; i < num_spatial_dims; ++i) {
         output_shape.push_back(input_shape[1 + i] / block_size_);
+        output_exprs.push_back(
+            (input_shape_with_exprs.expressions(1 + i) /
+             xla::DExpr::Const(block_size_))
+                .simplify());
       }
       output_shape.push_back(input_shape[feature_dim] * block_elems);
+      output_exprs.push_back(
+          (input_shape_with_exprs.expressions(feature_dim) *
+           xla::DExpr::Const(block_elems))
+              .simplify());
     } else {
       // FORMAT_NCHW
       int64_t block_elems = 1;
@@ -126,10 +147,17 @@ class SpaceToDepthOp : public XlaOpKernel {
       }
 
       reshaped_shape.push_back(input_shape[0]);
+      reshaped_exprs.push_back(input_shape_with_exprs.expressions(0));
       reshaped_shape.push_back(input_shape[feature_dim]);
+      reshaped_exprs.push_back(input_shape_with_exprs.expressions(feature_dim));
       for (int i = 0; i < num_spatial_dims; ++i) {
         reshaped_shape.push_back(input_shape[2 + i] / block_size_);
+        reshaped_exprs.push_back(
+            (input_shape_with_exprs.expressions(2 + i) /
+             xla::DExpr::Const(block_size_))
+                .simplify());
         reshaped_shape.push_back(block_size_);
+        reshaped_exprs.push_back(xla::DExpr::Const(block_size_));
       }
 
       transpose_order.push_back(0);
@@ -142,9 +170,18 @@ class SpaceToDepthOp : public XlaOpKernel {
       }
 
       output_shape.push_back(input_shape[0]);
+      output_exprs.push_back(input_shape_with_exprs.expressions(0));
       output_shape.push_back(input_shape[feature_dim] * block_elems);
+      output_exprs.push_back(
+          (input_shape_with_exprs.expressions(feature_dim) *
+           xla::DExpr::Const(block_elems))
+              .simplify());
       for (int i = 0; i < num_spatial_dims; ++i) {
         output_shape.push_back(input_shape[2 + i] / block_size_);
+        output_exprs.push_back(
+            (input_shape_with_exprs.expressions(2 + i) /
+             xla::DExpr::Const(block_size_))
+                .simplify());
       }
     }
 
@@ -156,7 +193,7 @@ class SpaceToDepthOp : public XlaOpKernel {
     //       input_shape[1] / block_size_, block_size_,
     //       input_shape[2] / block_size_, block_size_,
     //       depth]
-    xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape);
+    xla::XlaOp reshaped = xla::Reshape(input, reshaped_shape, reshaped_exprs);
 
     // 2. Permute dimensions of `reshaped` to produce
     //    `permuted_reshaped` of shape:
@@ -176,7 +213,7 @@ class SpaceToDepthOp : public XlaOpKernel {
     //       input_shape[2] / block_size_,
     //       block_size_ * block_size_ * depth]
     //
-    xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape);
+    xla::XlaOp output = xla::Reshape(permuted_reshaped, output_shape, output_exprs);
 
     // If this used to be a vectorized format turn it back now.
     if (data_format != data_format_) {
diff --git a/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc b/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc
index b4d589f183108e..709d5d8c28d17c 100644
--- a/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc
@@ -83,7 +83,8 @@ class SparseToDenseOp : public XlaOpKernel {
       sparse_values = Broadcast(sparse_values, {num_elems});
     }
     xla::XlaBuilder* builder = context->builder();
-    auto buffer = Broadcast(default_value, output_shape.dim_sizes());
+    auto buffer = Broadcast(default_value, output_shape.dim_sizes(),
+                            output_shape.get_filled_expressions());
     std::vector<bool> dynamic_dims;
     OP_REQUIRES_OK(
         context, context->ResolveInputDynamismIntoPredVector(1, &dynamic_dims));
diff --git a/tensorflow/compiler/tf2xla/kernels/split_op.cc b/tensorflow/compiler/tf2xla/kernels/split_op.cc
index 4f7c4ae99b6b6b..dcecb71e11e7d7 100644
--- a/tensorflow/compiler/tf2xla/kernels/split_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/split_op.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <cstdint>
 #include <vector>
 
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/xla_builder.h"
@@ -35,6 +36,8 @@ class SplitOp : public XlaOpKernel {
   explicit SplitOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
+    const bool enable_dynamic_sizes =
+        GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes;
     const int32_t num_split = num_outputs();
     const TensorShape split_dim_shape = ctx->InputShape("split_dim");
     const TensorShape input_shape = ctx->InputShape(1);
@@ -81,17 +84,25 @@ class SplitOp : public XlaOpKernel {
     // All the slices are the same size: this is the size along the
     // split dimension.
     const int32_t slice_size = input_shape.dim_size(split_dim) / num_split;
+    xla::DExpr slice_expr =
+        input_shape.get_filled_expression(split_dim) / num_split;
 
     // The vectors we will use to define the slice. The entry for the
     // split dimensions varies for each output.
     std::vector<int64_t> begin(input_shape.dims(), 0);
     std::vector<int64_t> limits(input_shape.dims());
+    std::vector<xla::DExpr> begin_expr;
+    std::vector<xla::DExpr> limits_expr;
+    begin_expr.reserve(input_shape.dims());
+    limits_expr.reserve(input_shape.dims());
     std::vector<int64_t> strides(input_shape.dims(), 1);
     for (int i = 0; i < input_shape.dims(); ++i) {
       // Initially set up the limits to be the full size of the input:
       // the split dimension is filled in below.
       int64_t dim = input_shape.dim_size(i);
       limits[i] = dim;
+      begin_expr.push_back(xla::DExpr::Const(0));
+      limits_expr.push_back(input_shape.get_filled_expression(i));
     }
 
     // Create each of the outputs.
@@ -99,7 +110,15 @@ class SplitOp : public XlaOpKernel {
       // Slice out the ith split from the split dimension.
       begin[split_dim] = i * slice_size;
       limits[split_dim] = (i + 1) * slice_size;
-      ctx->SetOutput(i, xla::Slice(input, begin, limits, strides));
+
+      if (enable_dynamic_sizes) {
+        begin_expr[split_dim] = (i * slice_expr).simplify();
+        limits_expr[split_dim] = ((i + 1) * slice_expr).simplify();
+        ctx->SetOutput(i, xla::Slice(input, begin, limits, begin_expr,
+                                     limits_expr, strides));
+      } else {
+        ctx->SetOutput(i, xla::Slice(input, begin, limits, strides));
+      }
     }
   }
 };
@@ -111,6 +130,8 @@ class SplitVOp : public XlaOpKernel {
   explicit SplitVOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
 
   void Compile(XlaOpKernelContext* ctx) override {
+    const bool enable_dynamic_sizes =
+        GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes;
     const int32_t num_split = num_outputs();
     const TensorShape input_shape = ctx->InputShape(0);
     const TensorShape index_shape = ctx->InputShape(2);
@@ -202,21 +223,38 @@ class SplitVOp : public XlaOpKernel {
           input_shape.dim_size(split_dim) - total_split_size;
     }
 
-    // The vectors we will use to define the slice. The entry for the
-    // split dimensions varies for each output.
+    // The vectors we will use to define the slice. The entry for the split
+    // dimension varies for each output.
     std::vector<int64_t> begin(input_shape.dims(), 0);
     auto dim_sizes = input_shape.dim_sizes();
     std::vector<int64_t> limits(dim_sizes.begin(), dim_sizes.end());
     std::vector<int64_t> strides(input_shape.dims(), 1);
+    std::vector<xla::DExpr> begin_expr(input_shape.dims(), xla::DExpr::Const(0));
+    auto input_exprs = input_shape.get_filled_expressions();
+    std::vector<xla::DExpr> limits_expr;
+    limits_expr.reserve(input_exprs.size());
+    for (auto expr : input_exprs) {
+      limits_expr.push_back(expr);
+    }
     for (int i = 0; i < num_split; ++i) {
-      TensorShape output_shape(input_shape);
       int slice_size = split_sizes[i];
-      output_shape.set_dim(split_dim, slice_size);
+      xla::DExpr slice_expr = xla::DExpr::Const(slice_size);
 
       // Slice out the ith split from the split dimension.
       limits[split_dim] = begin[split_dim] + slice_size;
-      ctx->SetOutput(i, xla::Slice(input, begin, limits, strides));
+      if (enable_dynamic_sizes) {
+        limits_expr[split_dim] =
+            (begin_expr[split_dim] + slice_expr).simplify();
+        ctx->SetOutput(i,
+                       xla::Slice(input, begin, limits, begin_expr,
+                                  limits_expr, strides));
+      } else {
+        ctx->SetOutput(i, xla::Slice(input, begin, limits, strides));
+      }
       begin[split_dim] = limits[split_dim];
+      if (enable_dynamic_sizes) {
+        begin_expr[split_dim] = limits_expr[split_dim];
+      }
     }
   }
 };
diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc
index 3c99ad63565266..28ca29002084c4 100644
--- a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc
@@ -157,7 +157,8 @@ class StackPushOp : public XlaOpKernel {
 
     TensorShape slice_shape = elem_shape;
     slice_shape.InsertDim(0, 1LL);
-    auto update = xla::Reshape(value, slice_shape.dim_sizes());
+    auto update = xla::Reshape(value, slice_shape.dim_sizes(),
+                               slice_shape.get_filled_expressions());
 
     // TODO(phawkins): We don't check the index is in bounds --- there is no
     // error mechanism in XLA.
diff --git a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
index aa71c5c34d2e1a..9e3d5b6fa6aad5 100644
--- a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc
@@ -49,8 +49,9 @@ xla::BitGeneratorTy GetBitGeneratorForDevice(
       device_type_string == DEVICE_CPU_XLA_JIT) {
     return [=](xla::XlaOp key, xla::XlaOp state, const xla::Shape& shape) {
       std::tie(state, key) = xla::ScramblePhiloxKey(key);
-      xla::XlaOp philox_state =
-          xla::ConcatInDim(key.builder(), {xla::Reshape(key, {1}), state}, 0);
+      xla::XlaOp philox_state = xla::ConcatInDim(
+          key.builder(),
+          {xla::Reshape(key, {1}, {xla::DExpr::Const(1)}), state}, 0);
       xla::XlaOp result = xla::RngBitGenerator(xla::RandomAlgorithm::RNG_PHILOX,
                                                philox_state, shape);
       return xla::RngOutput{/*value=*/xla::GetTupleElement(result, 1),
@@ -420,20 +421,23 @@ class StatelessParameterizedTruncatedNormalOp : public XlaOpKernel {
     auto rng_dtype = MaybeConvertBF16ToF32(dtype_);
     xla::Shape xla_shape;
     OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(rng_dtype, shape, &xla_shape));
-
-    auto bcasted_means = BroadcastTo(ctx->Input(2), shape.dim_sizes());
+    auto bcasted_means =
+        BroadcastTo(ctx->Input(2), shape.dim_sizes(), shape.get_filled_expressions());
     OP_REQUIRES_OK(ctx, bcasted_means.status());
     auto means = bcasted_means.value();
 
-    auto bcasted_stddevs = BroadcastTo(ctx->Input(3), shape.dim_sizes());
+    auto bcasted_stddevs =
+        BroadcastTo(ctx->Input(3), shape.dim_sizes(), shape.get_filled_expressions());
     OP_REQUIRES_OK(ctx, bcasted_stddevs.status());
     auto stddevs = bcasted_stddevs.value();
 
-    auto bcasted_minvals = BroadcastTo(ctx->Input(4), shape.dim_sizes());
+    auto bcasted_minvals =
+        BroadcastTo(ctx->Input(4), shape.dim_sizes(), shape.get_filled_expressions());
     OP_REQUIRES_OK(ctx, bcasted_minvals.status());
     auto minvals = bcasted_minvals.value();
 
-    auto bcasted_maxvals = BroadcastTo(ctx->Input(5), shape.dim_sizes());
+    auto bcasted_maxvals =
+        BroadcastTo(ctx->Input(5), shape.dim_sizes(), shape.get_filled_expressions());
     OP_REQUIRES_OK(ctx, bcasted_maxvals.status());
     auto maxvals = bcasted_maxvals.value();
 
diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
index e15196bd756462..8fd205097fe8ef 100644
--- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc
@@ -24,7 +24,9 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/types/span.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/tf2xla/literal_util.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
@@ -48,6 +50,37 @@ namespace tensorflow {
 namespace {
 using errors::InvalidArgument;
 
+bool TryBuildSlicedContents(const XlaExpression& input_expr,
+                            const TensorShape& input_shape,
+                            const absl::InlinedVector<int64_t, 4>& begin,
+                            const absl::InlinedVector<int64_t, 4>& strides,
+                            const TensorShape& final_shape,
+                            std::vector<xla::DExpr>* output_contents) {
+  output_contents->clear();
+  const auto& input_contents = input_expr.contents();
+  if (input_contents.empty() || input_shape.dims() != 1 || begin.size() != 1 ||
+      strides.size() != 1) {
+    return false;
+  }
+
+  const int64_t output_elements = final_shape.num_elements();
+  const int64_t start = begin[0];
+  const int64_t stride = strides[0];
+  for (int64_t i = 0; i < output_elements; ++i) {
+    const int64_t index = start + i * stride;
+    if (index < 0 || index >= input_contents.size()) {
+      output_contents->clear();
+      return false;
+    }
+    const xla::DExpr& expr = input_contents[index];
+    output_contents->push_back(
+        expr ? expr : xla::DExpr::Unknown(xla::kUnknownContentSentinel));
+  }
+  return absl::c_any_of(*output_contents, [](const xla::DExpr& expr) {
+    return expr && expr->is_dynamic();
+  });
+}
+
 class StridedSliceOp : public XlaOpKernel {
  public:
   explicit StridedSliceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
@@ -82,6 +115,9 @@ class StridedSliceOp : public XlaOpKernel {
         partial_final_shape.set_dim(
             i,
             input_shape.dim_size(shape_spec.output_to_processing_mapping[i]));
+        partial_final_shape.set_expression(
+            i, input_shape.get_expression(
+                   shape_spec.output_to_processing_mapping[i]));
       }
     }
 
@@ -97,6 +133,8 @@ class StridedSliceOp : public XlaOpKernel {
         // Use input shape to update unknown dimension of partial shape -- if a
         // dimension is unknown, we use input shape as bound.
         partial_processing_shape.set_dim(i, input_shape.dim_size(i));
+        partial_processing_shape.set_expression(i,
+                                                input_shape.get_expression(i));
       }
     }
     TensorShape processing_shape;
@@ -157,7 +195,10 @@ class StridedSliceOp : public XlaOpKernel {
       auto zero = xla::Zero(ctx->builder(), ctx->InputXlaType("begin"));
       xla::XlaOp begin_index, end_index;
       int64_t sparse_index = shape_spec.processing_to_sparse_mapping[i];
-      bool xla_input_is_dynamic = input_xla_shape.is_dynamic_dimension(i);
+      const xla::DExpr& input_expr = input_xla_shape.expressions(i);
+      bool xla_input_is_dynamic =
+          input_xla_shape.is_dynamic_dimension(i) ||
+          input_expr->is_dynamic();
       xla::XlaOp dim_size;
       if (xla_input_is_dynamic) {
         dim_size = xla::GetDimensionSize(ctx->Input(0), i);
@@ -215,10 +256,12 @@ class StridedSliceOp : public XlaOpKernel {
     }
 
     slice =
-        xla::DynamicSlice(slice, start_indices, processing_shape.dim_sizes());
+        xla::DynamicSlice(slice, start_indices, processing_shape.dim_sizes(),
+                          processing_shape.get_expressions());
     // new_axis_mask_, ellipsis_mask_ and shrink_axis_mask_ may add or remove
     // size 1 dims of a shape.
-    slice = xla::Reshape(slice, final_shape.dim_sizes());
+    slice = xla::Reshape(slice, final_shape.dim_sizes(),
+                         final_shape.get_expressions());
     for (int64_t i = 0; i < final_shape.dims(); ++i) {
       int64 processing_shape_dim = shape_spec.output_to_processing_mapping[i];
       // If processing_shape_dim is -1, it means the output dimension was newly
@@ -238,6 +281,8 @@ class StridedSliceOp : public XlaOpKernel {
   }
 
   void Compile(XlaOpKernelContext* ctx) override {
+    const bool enable_dynamic_sizes =
+        GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes;
     const TensorShape input_shape = ctx->InputShape(0);
     const TensorShape begin_shape = ctx->InputShape("begin");
     OP_REQUIRES(
@@ -246,6 +291,8 @@ class StridedSliceOp : public XlaOpKernel {
 
     absl::InlinedVector<int64_t, 4> begin;
     absl::InlinedVector<int64_t, 4> end;
+    absl::InlinedVector<xla::DExpr, 4> begin_expr;
+    absl::InlinedVector<xla::DExpr, 4> end_expr;
     absl::InlinedVector<int64_t, 4> strides;
 
     xla::Literal begin_literal, end_literal, strides_literal;
@@ -268,14 +315,15 @@ class StridedSliceOp : public XlaOpKernel {
     PartialTensorShape partial_processing_shape, partial_final_shape;
     bool dummy = false;
     StridedSliceShapeSpec shape_spec;
+
     OP_REQUIRES_OK(
-        ctx,
-        ValidateStridedSliceOp(
-            begin_is_constant ? &begin_tensor : nullptr,
-            end_is_constant ? &end_tensor : nullptr, strides_tensor,
-            input_shape, begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_,
-            shrink_axis_mask_, &partial_processing_shape, &partial_final_shape,
-            &dummy, &dummy, &dummy, &begin, &end, &strides, &shape_spec));
+        ctx, ValidateStridedSliceOp(
+                 begin_is_constant ? &begin_tensor : nullptr,
+                 end_is_constant ? &end_tensor : nullptr, strides_tensor,
+                 input_shape, begin_mask_, end_mask_, ellipsis_mask_,
+                 new_axis_mask_, shrink_axis_mask_, &partial_processing_shape,
+                 &partial_final_shape, &dummy, &dummy, &dummy, &begin, &end,
+                 &strides, &begin_expr, &end_expr, &shape_spec));
 
     xla::XlaOp slice = ctx->Input(0);
     std::vector<bool> begins_are_dynamic;
@@ -294,17 +342,30 @@ class StridedSliceOp : public XlaOpKernel {
                           ", output shape must be a compile-time constant"));
       absl::InlinedVector<int64_t, 4> dimensions_to_reverse;
       absl::InlinedVector<int64_t, 4> slice_begin, slice_end, slice_strides;
+      absl::InlinedVector<xla::DExpr, 4> slice_begin_expr, slice_end_expr;
       for (int i = 0; i < begin.size(); ++i) {
         if (strides[i] > 0) {
           slice_begin.push_back(begin[i]);
+          slice_begin_expr.push_back(begin_expr[i]);
           slice_end.push_back(std::max(end[i], begin[i]));
+          slice_end_expr.push_back(
+              xla::DExpr::Max(end_expr[i], begin_expr[i]).simplify());
           slice_strides.push_back(strides[i]);
         } else {
           // Negative stride: swap begin and end, add 1 because the interval
           // is semi-open, and mark the dimension to be reversed.
+          xla::DExpr input_expr = input_shape.get_filled_expression(i);
           slice_begin.push_back(input_shape.dim_size(i) - begin[i] - 1);
+          slice_begin_expr.push_back(
+              (input_expr - begin_expr[i] - xla::DExpr::Const(1)).simplify());
           slice_end.push_back(std::max(input_shape.dim_size(i) - end[i] - 1,
                                        input_shape.dim_size(i) - begin[i] - 1));
+          slice_end_expr.push_back(
+              xla::DExpr::Max(
+                  (input_expr - end_expr[i] - xla::DExpr::Const(1)).simplify(),
+                  (input_expr - begin_expr[i] - xla::DExpr::Const(1))
+                      .simplify())
+                  .simplify());
           slice_strides.push_back(-strides[i]);
           dimensions_to_reverse.push_back(i);
         }
@@ -312,7 +373,40 @@ class StridedSliceOp : public XlaOpKernel {
       if (!dimensions_to_reverse.empty()) {
         slice = xla::Rev(slice, dimensions_to_reverse);
       }
-      slice = xla::Slice(slice, slice_begin, slice_end, slice_strides);
+      for (int i = 0; i < partial_processing_shape.dims(); ++i) {
+        partial_processing_shape.set_expression(
+            i, ((slice_end_expr[i] - slice_begin_expr[i] +
+                 xla::DExpr::Const(slice_strides[i]) - xla::DExpr::Const(1)) /
+                xla::DExpr::Const(slice_strides[i]))
+                   .simplify());
+      }
+      for (int i = 0; i < partial_final_shape.dims(); ++i) {
+        int64_t processing_index = shape_spec.output_to_processing_mapping[i];
+        partial_final_shape.set_expression(
+            i, processing_index == -1
+                   ? xla::DExpr::Const(partial_final_shape.dim_size(i))
+                   : partial_processing_shape.get_filled_expression(
+                         processing_index));
+      }
+      OP_REQUIRES(
+          ctx, partial_final_shape.AsTensorShape(&final_shape),
+          InvalidArgument("XLA can't deduce compile time constant output "
+                          "shape for strided slice: ",
+                          partial_final_shape.DebugString(),
+                          ", output shape must be a compile-time constant"));
+      slice = enable_dynamic_sizes
+                  ? xla::Slice(slice, slice_begin, slice_end, slice_begin_expr,
+                               slice_end_expr, slice_strides)
+                  : xla::Slice(slice, slice_begin, slice_end, slice_strides);
+      std::vector<xla::DExpr> output_contents;
+      const bool has_output_contents =
+          SymbolicContentEnabled() &&
+          TryBuildSlicedContents(ctx->InputExpression(0), input_shape, begin,
+                                 strides, final_shape, &output_contents);
+      if (has_output_contents) {
+        OP_REQUIRES_OK(
+            ctx, ctx->builder()->SetInstructionContents(slice, output_contents));
+      }
       auto operand_shape_or = ctx->builder()->GetShape(ctx->Input(0));
       OP_REQUIRES_OK(ctx, operand_shape_or.status());
       xla::Shape xla_shape = operand_shape_or.value();
@@ -325,8 +419,20 @@ class StridedSliceOp : public XlaOpKernel {
       bool ends_are_static = absl::c_all_of(
           ends_are_dynamic, [](bool dynamic) { return !dynamic; });
       // Static output shape, return a static slice.
-      slice = xla::Reshape(slice, final_shape.dim_sizes());
+      slice = xla::Reshape(slice, final_shape.dim_sizes(),
+                           final_shape.get_expressions());
+      if (has_output_contents) {
+        OP_REQUIRES_OK(
+            ctx, ctx->builder()->SetInstructionContents(slice, output_contents));
+      }
       if (xla_shape.is_static() && ends_are_static) {
+        if (has_output_contents) {
+          auto output_expr =
+              XlaExpression::XlaOp(slice, ctx->expected_output_dtype(0));
+          output_expr.set_contents(std::move(output_contents));
+          ctx->SetOutputExpression(0, output_expr);
+          return;
+        }
         ctx->SetOutput(0, slice);
         return;
       }
@@ -385,8 +491,20 @@ class StridedSliceOp : public XlaOpKernel {
               xla::Sub(operand_size, xla::ConstantR0<int32>(
                                          ctx->builder(), begin[input_index])),
               i);
+          if (has_output_contents) {
+            OP_REQUIRES_OK(
+                ctx, ctx->builder()->SetInstructionContents(slice,
+                                                            output_contents));
+          }
         }
       }
+      if (has_output_contents) {
+        auto output_expr =
+            XlaExpression::XlaOp(slice, ctx->expected_output_dtype(0));
+        output_expr.set_contents(std::move(output_contents));
+        ctx->SetOutputExpression(0, output_expr);
+        return;
+      }
       ctx->SetOutput(0, slice);
       return;
     } else {
@@ -431,11 +549,15 @@ class StridedSliceGradOp : public XlaOpKernel {
   void CompileAsDynamicUpdateSlice(XlaOpKernelContext* ctx,
                                    const TensorShape& input_shape,
                                    const xla::Literal& strides_literal) {
+    const bool enable_dynamic_sizes =
+        GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes;
     bool dummy = false;
     Tensor strides_tensor;
     PartialTensorShape processing_shape, final_shape;
     absl::InlinedVector<int64_t, 4> begin;
     absl::InlinedVector<int64_t, 4> end;
+    absl::InlinedVector<xla::DExpr, 4> begin_expr;
+    absl::InlinedVector<xla::DExpr, 4> end_expr;
     absl::InlinedVector<int64_t, 4> strides;
     StridedSliceShapeSpec shape_spec;
     OP_REQUIRES_OK(ctx, LiteralToHostTensor(strides_literal, index_type_,
@@ -445,7 +567,7 @@ class StridedSliceGradOp : public XlaOpKernel {
                  nullptr, nullptr, strides_tensor, input_shape, begin_mask_,
                  end_mask_, ellipsis_mask_, new_axis_mask_, shrink_axis_mask_,
                  &processing_shape, &final_shape, &dummy, &dummy, &dummy,
-                 &begin, &end, &strides, &shape_spec));
+                 &begin, &end, &strides, &begin_expr, &end_expr, &shape_spec));
     for (int64_t i = 0; i < processing_shape.dims(); ++i) {
       OP_REQUIRES(
           ctx, strides[i] == 1,
@@ -459,14 +581,22 @@ class StridedSliceGradOp : public XlaOpKernel {
     VLOG(1) << "xla final_shape" << final_shape;
     VLOG(1) << "input_shape" << input_shape.DebugString();
     auto input_sizes = input_shape.dim_sizes();
+    std::vector<xla::DExpr> input_exprs;
+    input_exprs.reserve(input_shape.dims());
+    for (int64_t i = 0; i < input_shape.dims(); ++i) {
+      input_exprs.push_back(input_shape.get_filled_expression(i));
+    }
     // For unknown output dim the bound of the output shape is input.  Pad and
     // double the size of input shape to leave enough buffer to avoid OOB
     // dynamic update slice.
     auto input_sizes_padded = input_shape.dim_sizes();
+    auto input_exprs_padded = input_exprs;
     bool need_padding = false;
     for (int64_t i = 0; i < processing_shape.dims(); ++i) {
       if (processing_shape.dim_size(i) == -1) {
         input_sizes_padded[i] *= 2;
+        input_exprs_padded[i] =
+            (xla::DExpr::Const(2) * input_exprs_padded[i]).simplify();
         need_padding = true;
       }
     }
@@ -477,6 +607,7 @@ class StridedSliceGradOp : public XlaOpKernel {
       if (shape_spec.output_to_processing_mapping[i] != -1) {
         processing_shape.set_dim(shape_spec.output_to_processing_mapping[i],
                                  grad_shape.dimensions(i));
+        // TODO Pass it back
       }
     }
 
@@ -506,15 +637,21 @@ class StridedSliceGradOp : public XlaOpKernel {
     }
     auto zero = XlaHelpers::Zero(ctx->builder(), ctx->expected_output_dtype(0));
 
-    zero = xla::Broadcast(zero, input_sizes_padded);
-    grad = xla::Reshape(grad, processing_shape.dim_sizes());
+    zero = xla::Broadcast(zero, input_sizes_padded, input_exprs_padded);
+    grad = xla::Reshape(grad, processing_shape.dim_sizes(),
+                        processing_shape.get_expressions());
     grad = xla::DynamicUpdateSlice(zero, grad, begins);
     if (need_padding) {
       // We padded the input shape to avoid OOB when DUS. Now slice out the
       // padding in the final result.
       std::vector<int64_t> strides(input_shape.dims(), 1);
       std::vector<int64_t> start_indices(input_shape.dims(), 0);
-      grad = xla::Slice(grad, start_indices, input_sizes, strides);
+      std::vector<xla::DExpr> start_exprs(input_shape.dims(),
+                                          xla::DExpr::Const(0));
+      grad = enable_dynamic_sizes
+                 ? xla::Slice(grad, start_indices, input_sizes, start_exprs,
+                              input_exprs, strides)
+                 : xla::Slice(grad, start_indices, input_sizes, strides);
     }
     ctx->SetOutput(0, grad);
   }
@@ -522,6 +659,8 @@ class StridedSliceGradOp : public XlaOpKernel {
     TensorShape processing_shape, final_shape;
     absl::InlinedVector<int64_t, 4> begin;
     absl::InlinedVector<int64_t, 4> end;
+    absl::InlinedVector<xla::DExpr, 4> begin_expr;
+    absl::InlinedVector<xla::DExpr, 4> end_expr;
     absl::InlinedVector<int64_t, 4> strides;
 
     TensorShape input_shape;
@@ -547,11 +686,13 @@ class StridedSliceGradOp : public XlaOpKernel {
 
     bool dummy = false;
     OP_REQUIRES_OK(
-        ctx, ValidateStridedSliceOp(
-                 &begin_tensor, &end_tensor, strides_tensor, input_shape,
-                 begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_,
-                 shrink_axis_mask_, &processing_shape, &final_shape, &dummy,
-                 &dummy, &dummy, &begin, &end, &strides));
+        ctx,
+        ValidateStridedSliceOp(
+            &begin_tensor, &end_tensor, strides_tensor, input_shape,
+            begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_,
+            shrink_axis_mask_, &processing_shape, &final_shape, &dummy, &dummy,
+            &dummy, &begin, &end, &strides, &begin_expr, &end_expr,
+            /*shape_spec=*/nullptr));
 
     // Check to make sure dy is consistent with the original slice
     const TensorShape dy_shape = ctx->InputShape(4);
@@ -570,7 +711,8 @@ class StridedSliceGradOp : public XlaOpKernel {
     xla::XlaOp grad = ctx->Input(4);
 
     // Undo any new/shrink axes.
-    grad = xla::Reshape(grad, processing_shape.dim_sizes());
+    grad = xla::Reshape(grad, processing_shape.dim_sizes(),
+                        processing_shape.get_expressions());
 
     // Pad the input gradients.
     absl::InlinedVector<int64_t, 4> dimensions_to_reverse;
@@ -662,6 +804,8 @@ class StridedSliceAssignOp : public XlaOpKernel {
     TensorShape final_shape;
     absl::InlinedVector<int64_t, 4> begin;
     absl::InlinedVector<int64_t, 4> end;
+    absl::InlinedVector<xla::DExpr, 4> begin_expr;
+    absl::InlinedVector<xla::DExpr, 4> end_expr;
     absl::InlinedVector<int64_t, 4> strides;
 
     xla::Literal begin_literal, end_literal, strides_literal;
@@ -690,12 +834,14 @@ class StridedSliceAssignOp : public XlaOpKernel {
 
     TensorShape dummy_processing_shape;
     bool dummy = false;
-    OP_REQUIRES_OK(ctx,
-                   ValidateStridedSliceOp(
-                       &begin_tensor, &end_tensor, strides_tensor, lhs_shape,
-                       begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_,
-                       shrink_axis_mask_, &dummy_processing_shape, &final_shape,
-                       &dummy, &dummy, &dummy, &begin, &end, &strides));
+    OP_REQUIRES_OK(
+        ctx,
+        ValidateStridedSliceOp(
+            &begin_tensor, &end_tensor, strides_tensor, lhs_shape, begin_mask_,
+            end_mask_, ellipsis_mask_, new_axis_mask_, shrink_axis_mask_,
+            &dummy_processing_shape, &final_shape, &dummy, &dummy, &dummy,
+            &begin, &end, &strides, &begin_expr, &end_expr,
+            /*shape_spec=*/nullptr));
 
     if (final_shape.num_elements() == 0 && rhs_shape.num_elements() == 0) {
       // DynamicUpdateSlice does not allow 0-element updates. We should probably
@@ -717,6 +863,7 @@ class StridedSliceAssignOp : public XlaOpKernel {
     absl::InlinedVector<int64_t, 4> dimensions_to_reverse;
     absl::InlinedVector<xla::XlaOp, 4> slice_begin;
     absl::InlinedVector<int64_t, 4> slice_dims;
+    absl::InlinedVector<xla::DExpr, 4> slice_exprs;
     for (int i = 0; i < begin.size(); ++i) {
       // TODO(b/121179231): implement strides != 1
       OP_REQUIRES(
@@ -726,12 +873,14 @@ class StridedSliceAssignOp : public XlaOpKernel {
         slice_begin.push_back(
             xla::ConstantR0<int64_t>(ctx->builder(), begin[i]));
         slice_dims.push_back(end[i] - begin[i]);
+        slice_exprs.push_back(xla::DExpr::Const(end[i] - begin[i]));
       } else {
         // Negative stride: swap begin and end, add 1 because the interval
         // is semi-open, and mark the dimension to be reversed.
         slice_begin.push_back(
             xla::ConstantR0<int64_t>(ctx->builder(), end[i] + 1));
         slice_dims.push_back(begin[i] - end[i]);
+        slice_exprs.push_back(xla::DExpr::Const(begin[i] - end[i]));
         dimensions_to_reverse.push_back(i);
       }
     }
@@ -739,7 +888,7 @@ class StridedSliceAssignOp : public XlaOpKernel {
     if (!dimensions_to_reverse.empty()) {
       rhs = xla::Rev(rhs, dimensions_to_reverse);
     }
-    rhs = xla::Reshape(rhs, slice_dims);
+    rhs = xla::Reshape(rhs, slice_dims, slice_exprs);
 
     lhs = xla::DynamicUpdateSlice(lhs, rhs, slice_begin);
 
diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc
index 888908e30b2331..9bc5d0d95bfd8c 100644
--- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc
@@ -165,9 +165,11 @@ class TensorArrayOp : public XlaOpKernel {
       CHECK(element_shape_.AsTensorShape(&shape));
       TensorShape ta_shape;
       ta_shape.AddDim(size);
+      ta_shape.AddExpression(xla::DExpr::Const(size));
       ta_shape.AppendShape(shape);
       xla::XlaOp zero = XlaHelpers::Zero(b, dtype_);
-      value = xla::Broadcast(zero, ta_shape.dim_sizes());
+      value = xla::Broadcast(zero, ta_shape.dim_sizes(),
+                             ta_shape.get_filled_expressions());
     }
 
     XlaResource* var =
@@ -223,7 +225,8 @@ class TensorArrayWriteOp : public XlaOpKernel {
 
     TensorShape slice_shape = elem_shape;
     slice_shape.InsertDim(0, 1LL);
-    auto update = xla::Reshape(value, slice_shape.dim_sizes());
+    auto update = xla::Reshape(value, slice_shape.dim_sizes(),
+                               slice_shape.get_filled_expressions());
 
     xla::XlaOp written;
     if (resource->tensor_array_multiple_writes_aggregate()) {
@@ -274,9 +277,12 @@ class TensorArrayReadOp : public XlaOpKernel {
     start_indices[0] = index;
 
     auto slice_shape = ta_shape.dim_sizes();
+    auto slice_exprs = ta_shape.get_filled_expressions();
     slice_shape[0] = 1LL;
+    slice_exprs[0] = xla::DExpr::Const(1LL);
 
-    xla::XlaOp read = xla::DynamicSlice(ta, start_indices, slice_shape);
+    xla::XlaOp read =
+        xla::DynamicSlice(ta, start_indices, slice_shape, slice_exprs);
 
     // Remove the leading '1' dimension.
     std::vector<int64_t> value_shape(slice_shape.begin() + 1,
@@ -469,9 +475,16 @@ class TensorArrayConcatOp : public XlaOpKernel {
     xla::XlaOp ta = resource->value();
 
     auto ta_dims = ta_shape.dim_sizes();
+    auto ta_exprs = ta_shape.get_filled_expressions();
     std::vector<int64_t> shape(ta_dims.begin() + 1, ta_dims.end());
+    std::vector<xla::DExpr> exprs;
+    exprs.reserve(ta_exprs.size() - 1);
+    for (auto it = ta_exprs.begin() + 1; it != ta_exprs.end(); ++it) {
+      exprs.push_back(*it);
+    }
     shape[0] *= ta_shape.dim_size(0);
-    ctx->SetOutput(0, xla::Reshape(ta, shape));
+    exprs[0] = ta_exprs[0] * ta_shape.get_filled_expression(0);
+    ctx->SetOutput(0, xla::Reshape(ta, shape, exprs));
 
     Tensor lengths(DT_INT64, {ta_dims[0]});
     auto lengths_vec = lengths.vec<int64_t>();
@@ -526,6 +539,7 @@ class TensorArraySplitOp : public XlaOpKernel {
 
     TensorShape ta_shape;
     ta_shape.AddDim(resource->max_array_size());
+    ta_shape.AddExpression(xla::DExpr::Const(resource->max_array_size()));
     ta_shape.AppendShape(elem_shape);
 
     OP_REQUIRES(ctx, lengths.size() == resource->max_array_size(),
@@ -541,7 +555,8 @@ class TensorArraySplitOp : public XlaOpKernel {
                                         value_shape.DebugString(), " vs. ",
                                         ta_shape.DebugString()));
 
-    const xla::XlaOp reshape = xla::Reshape(value, ta_shape.dim_sizes());
+    const xla::XlaOp reshape =
+        xla::Reshape(value, ta_shape.dim_sizes(), ta_shape.get_filled_expressions());
     if (dtype_ == DT_BOOL) {
       ta = xla::Or(ta, reshape);
     } else {
diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc
index a1f58d5ae9b40e..92128c0dc19873 100644
--- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc
@@ -131,7 +131,8 @@ absl::Status TryGetElementShapeFromInput(XlaOpKernelContext* ctx,
     return absl::OkStatus();
   }
 
-  *shape = xla::ShapeUtil::MakeShape(dtype, partial_shape.dim_sizes());
+  *shape = xla::ShapeUtil::MakeShape(dtype, partial_shape.dim_sizes(),
+                                     partial_shape.get_filled_expressions());
   *got_shape = true;
   return absl::OkStatus();
 }
@@ -503,6 +504,8 @@ class TensorListConcatOp : public XlaOpKernel {
     xla::Shape element_shape = std::move(shape_or).value();
     std::vector<int64_t> element_dims =
         xla::SpanToVector(element_shape.dimensions());
+    std::vector<xla::DExpr> element_exprs(element_shape.expressions().begin(),
+                                          element_shape.expressions().end());
     OP_REQUIRES(
         ctx, element_dims.size() > 1,
         errors::Unimplemented("TensorList of scalars is not supported"));
@@ -510,12 +513,15 @@ class TensorListConcatOp : public XlaOpKernel {
     int64_t tensor_lengths = element_dims[1];
 
     std::vector<int64_t> new_dims = {num_elements * tensor_lengths};
+    std::vector<xla::DExpr> new_exprs = {
+        xla::DExpr::Const(num_elements * tensor_lengths)};
 
     for (int i = 2; i < element_dims.size(); i++) {
       new_dims.push_back(element_dims[i]);
+      new_exprs.push_back(element_exprs[i]);
     }
 
-    xla::XlaOp out = xla::Reshape(buffer, new_dims);
+    xla::XlaOp out = xla::Reshape(buffer, new_dims, new_exprs);
     ctx->SetOutput(0, out);
 
     // Second output is a tensor of lengths of returned tensors.
@@ -550,6 +556,8 @@ class TensorListSplitOp : public XlaOpKernel {
     xla::Shape element_shape = std::move(shape_or).value();
     std::vector<int64_t> element_dims =
         xla::SpanToVector(element_shape.dimensions());
+    std::vector<xla::DExpr> element_exprs(element_shape.expressions().begin(),
+                                          element_shape.expressions().end());
     OP_REQUIRES(
         ctx, !element_dims.empty(),
         errors::Unimplemented("Element dimensions have to be non-empty"));
@@ -569,11 +577,13 @@ class TensorListSplitOp : public XlaOpKernel {
         ctx, element_dims[0] % length == 0,
         errors::Unimplemented("Buffer size has to be a multiple of length"));
     std::vector<int64_t> new_dims = {element_dims[0] / length, length};
+    std::vector<xla::DExpr> new_exprs = {element_exprs[0] / xla::DExpr::Const(length),
+                                         xla::DExpr::Const(length)};
     for (int i = 1; i < element_dims.size(); i++) {
       new_dims.push_back(element_dims[i]);
     }
 
-    xla::XlaOp reshaped = xla::Reshape(input_tensor, new_dims);
+    xla::XlaOp reshaped = xla::Reshape(input_tensor, new_dims, new_exprs);
 
     xla::XlaOp result;
     OP_REQUIRES_OK(ctx, ExecuteTensorListFromTensor(length, reshaped, &result));
diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc
index 683dc4737e6dab..412c01f24dfc19 100644
--- a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc
+++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc
@@ -241,9 +241,12 @@ absl::Status GetTensorListShapeFromElementTensorListShape(
     const xla::Shape& shape =
         xla::ShapeUtil::GetTupleElementShape(element_tensor_list_shape, i);
     std::vector<int64_t> dimensions = xla::SpanToVector(shape.dimensions());
+    std::vector<xla::DExpr> expressions = xla::SpanToVector(shape.expressions());
     dimensions.insert(dimensions.begin(), leading_dim);
+    expressions.insert(expressions.begin(), xla::DExpr::Const(leading_dim));
     shapes.push_back(
-        xla::ShapeUtil::MakeShape(shape.element_type(), dimensions));
+        xla::ShapeUtil::MakeShape(shape.element_type(), dimensions,
+                                  expressions));
     if (leading_dim_is_dynamic) {
       shapes.back().set_dynamic_dimension(0, true);
     }
@@ -267,9 +270,13 @@ absl::Status GetTensorListShapeFromElementShape(const xla::Shape& element_shape,
   std::vector<xla::Shape> shapes;
   std::vector<int64_t> dimensions =
       xla::SpanToVector(element_shape.dimensions());
+  std::vector<xla::DExpr> expressions =
+      xla::SpanToVector(element_shape.expressions());
   dimensions.insert(dimensions.begin(), leading_dim);
+  expressions.insert(expressions.begin(), xla::DExpr::Const(leading_dim));
   shapes.push_back(
-      xla::ShapeUtil::MakeShape(element_shape.element_type(), dimensions));
+      xla::ShapeUtil::MakeShape(element_shape.element_type(), dimensions,
+                                expressions));
   shapes.back().set_dynamic_dimension(0, leading_dim_is_dynamic);
   shapes.push_back(xla::ShapeUtil::MakeShape(xla::PrimitiveType::S32,
                                              std::vector<int64_t>{}));
@@ -289,7 +296,8 @@ absl::Status CreateZerosTensorListWithShape(
         xla::ShapeUtil::GetTupleElementShape(list_shape, i);
     xla::XlaOp zero =
         xla::ConstantLiteral(b, xla::LiteralUtil::Zero(shape.element_type()));
-    xla::XlaOp zeros = xla::Broadcast(zero, shape.dimensions());
+    xla::XlaOp zeros =
+        xla::Broadcast(zero, shape.dimensions(), shape.expressions());
     TF_RET_CHECK(dynamic_dims[i].size() == shape.dimensions().size());
     for (int64_t dim = 0; dim < shape.dimensions().size(); ++dim) {
       if (shape.is_dynamic_dimension(dim)) {
@@ -389,7 +397,14 @@ absl::Status ExecuteTensorListPushBack(xla::XlaOp list, xla::XlaOp element,
       std::vector<int64_t> element_part_dims =
           xla::SpanToVector(element_part_shape.dimensions());
       element_part_dims.insert(element_part_dims.begin(), 1);
-      element_part = xla::Reshape(element_part, element_part_dims);
+      std::vector<xla::DExpr> element_part_exprs;
+      element_part_exprs.reserve(element_part_shape.expressions().size() + 1);
+      element_part_exprs.push_back(xla::DExpr::Const(1));
+      for (auto expr : element_part_shape.expressions()) {
+        element_part_exprs.push_back(expr);
+      }
+      element_part =
+          xla::Reshape(element_part, element_part_dims, element_part_exprs);
 
       std::vector<xla::XlaOp> start_indices(
           element_part_shape.dimensions().size() + 1,
@@ -406,7 +421,13 @@ absl::Status ExecuteTensorListPushBack(xla::XlaOp list, xla::XlaOp element,
     std::vector<int64_t> element_dims =
         xla::SpanToVector(element_shape.dimensions());
     element_dims.insert(element_dims.begin(), 1);
-    xla::XlaOp update = xla::Reshape(element, element_dims);
+    std::vector<xla::DExpr> element_exprs;
+    element_exprs.reserve(element_shape.expressions().size() + 1);
+    element_exprs.push_back(xla::DExpr::Const(1));
+    for (auto expr : element_shape.expressions()) {
+      element_exprs.push_back(expr);
+    }
+    xla::XlaOp update = xla::Reshape(element, element_dims, element_exprs);
 
     std::vector<xla::XlaOp> start_indices(element_shape.dimensions().size() + 1,
                                           xla::ConstantR0<int32>(b, 0));
@@ -455,11 +476,19 @@ absl::Status ExecuteTensorListPopBack(xla::XlaOp list, xla::XlaOp* list_result,
         xla::SpanToVector(list_part_shape.dimensions());
     slice_shape[0] = 1LL;
 
+    std::vector<xla::DExpr> slice_exprs;
+    slice_exprs.reserve(list_part_shape.expressions().size());
+    for (auto expr : list_part_shape.expressions()) {
+      slice_exprs.push_back(expr);
+    }
+    slice_exprs[0] = xla::DExpr::Const(1LL);
+
     xla::XlaOp list_part = xla::GetTupleElement(list, i);
     xla::XlaOp read = xla::DynamicSlice(list_part, start_indices, slice_shape);
 
     slice_shape.erase(slice_shape.begin());
-    element_result_parts.push_back(xla::Reshape(read, slice_shape));
+    slice_exprs.erase(slice_exprs.begin());
+    element_result_parts.push_back(xla::Reshape(read, slice_shape, slice_exprs));
     list_result_parts.push_back(list_part);
   }
   list_result_parts.push_back(push_index);
@@ -493,7 +522,13 @@ absl::Status ExecuteTensorListSetItem(xla::XlaOp list, xla::XlaOp index,
   std::vector<int64_t> element_dims =
       xla::SpanToVector(element_shape.dimensions());
   element_dims.insert(element_dims.begin(), 1);
-  xla::XlaOp update = xla::Reshape(element, element_dims);
+  std::vector<xla::DExpr> element_exprs;
+  element_exprs.reserve(element_shape.expressions().size() + 1);
+  element_exprs.push_back(xla::DExpr::Const(1));
+  for (auto expr : element_shape.expressions()) {
+    element_exprs.push_back(expr);
+  }
+  xla::XlaOp update = xla::Reshape(element, element_dims, element_exprs);
 
   std::vector<xla::XlaOp> start_indices(element_shape.dimensions().size() + 1,
                                         xla::ConstantR0<int32>(b, 0));
@@ -557,6 +592,13 @@ absl::Status ExecuteTensorListGetItem(xla::XlaOp list, xla::XlaOp index,
       xla::SpanToVector(buffer_shape.dimensions());
   slice_shape[0] = 1LL;
 
+  std::vector<xla::DExpr> slice_exprs;
+  slice_exprs.reserve(buffer_shape.expressions().size());
+  for (auto expr : buffer_shape.expressions()) {
+    slice_exprs.push_back(expr);
+  }
+  slice_exprs[0] = xla::DExpr::Const(1LL);
+
   xla::XlaOp list_part = xla::GetTupleElement(list, 0);
   xla::XlaOp read = xla::DynamicSlice(list_part, start_indices, slice_shape);
   // Propagate dynamic dimensions from buffer to the sliced buffer, except for
@@ -569,7 +611,8 @@ absl::Status ExecuteTensorListGetItem(xla::XlaOp list, xla::XlaOp index,
     }
   }
   slice_shape.erase(slice_shape.begin());
-  *result = xla::Reshape(read, slice_shape);
+  slice_exprs.erase(slice_exprs.begin());
+  *result = xla::Reshape(read, slice_shape, slice_exprs);
   return absl::OkStatus();
 }
 
diff --git a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc
index 6c39981ba5b937..10752ebb0f8c25 100644
--- a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc
@@ -67,17 +67,35 @@ class TileOp : public XlaOpKernel {
                             xla::ValueInferenceMode::kUpperBound));
 
     std::vector<int64_t> output_dims(input_shape.dims());
+    std::vector<xla::DExpr> output_exprs(input_shape.dims());
+
+    auto expr_sizes = input_shape.get_filled_expressions();
+    const auto& multiples_contents =
+        ctx->InputExpression("multiples").contents();
+
     for (int64_t i = 0; i < input_shape.dims(); ++i) {
       OP_REQUIRES(ctx, multiples_bounds[i] >= 0,
                   errors::InvalidArgument("Expected multiples[", i,
                                           "] >= 0, but got ", output_dims[i]));
       output_dims[i] = input_shape.dim_size(i) * multiples_bounds[i];
+      const xla::DExpr multiple_expr =
+          i < multiples_contents.size() && multiples_contents[i] &&
+                  multiples_contents[i]->is_dynamic()
+              ? multiples_contents[i]
+              : xla::DExpr::Const(multiples_bounds[i]);
+      output_exprs[i] = (expr_sizes[i] * multiple_expr).simplify();
     }
 
     std::vector<bool> multiples_are_dynamic;
 
     OP_REQUIRES_OK(ctx, ctx->ResolveInputDynamismIntoPredVector(
                             1, &multiples_are_dynamic));
+    for (int64_t i = 0; i < multiples_are_dynamic.size(); ++i) {
+      if (i < multiples_contents.size() && multiples_contents[i] &&
+          multiples_contents[i]->is_dynamic()) {
+        multiples_are_dynamic[i] = true;
+      }
+    }
 
     bool all_multiples_are_static = absl::c_all_of(
         multiples_are_dynamic, [](bool dynamic) { return !dynamic; });
@@ -91,8 +109,8 @@ class TileOp : public XlaOpKernel {
         return;
       }
     }
-
-    auto result_or = BroadcastTo(ctx->Input("input"), output_dims);
+    auto result_or =
+        BroadcastTo(ctx->Input("input"), output_dims, output_exprs);
 
     OP_REQUIRES_OK(ctx, result_or.status());
     auto result = result_or.value();
diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
index f35e375356516b..8dc75be96d7958 100644
--- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
@@ -45,6 +45,22 @@ namespace {
   };                                                                   \
   REGISTER_XLA_OP(Name(#NAME), NAME##Op);
 
+#define XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(NAME, COMPUTATION)        \
+  class NAME##NativeOp : public XlaOpKernel {                          \
+   public:                                                             \
+    explicit NAME##NativeOp(OpKernelConstruction* ctx)                 \
+        : XlaOpKernel(ctx) {}                                          \
+    void Compile(XlaOpKernelContext* ctx) override {                   \
+      xla::XlaBuilder* b = ctx->builder();                             \
+      (void)b;                                                         \
+      xla::XlaOp x = ctx->Input(0);                                    \
+      xla::XlaOp y = COMPUTATION;                                      \
+      ctx->SetOutput(0, y);                                            \
+    }                                                                  \
+  };                                                                   \
+  REGISTER_XLA_OP_FACTORY(                                             \
+      Name(#NAME), CreateDynamicNativeXlaOpKernel<NAME##NativeOp>);
+
 XLAJIT_MAKE_UNARY(ComplexAbs, xla::Abs(x));
 
 XLAJIT_MAKE_UNARY(Angle, xla::Atan2(xla::Imag(x), xla::Real(x)));
@@ -52,29 +68,30 @@ XLAJIT_MAKE_UNARY(Angle, xla::Atan2(xla::Imag(x), xla::Real(x)));
 XLAJIT_MAKE_UNARY(Conj, xla::Conj(x));
 
 // Return x if x>0, otherwise -x.
-REGISTER_XLA_OP(Name("Abs"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Abs, xla::Abs(x));
 XLAJIT_MAKE_UNARY(Acos, xla::Acos(x));
 XLAJIT_MAKE_UNARY(Acosh, xla::Acosh(x));
 XLAJIT_MAKE_UNARY(Asin, xla::Asin(x))
 XLAJIT_MAKE_UNARY(Asinh, xla::Asinh(x));
-REGISTER_XLA_OP(Name("Atan"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Atan, xla::Atan(x));
 XLAJIT_MAKE_UNARY(Atanh, xla::Atanh(x));
-REGISTER_XLA_OP(Name("Ceil"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("Cos"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Ceil, xla::Ceil(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Cos, xla::Cos(x));
 XLAJIT_MAKE_UNARY(Cosh, xla::Cosh(x));
 XLAJIT_MAKE_UNARY(Sin, xla::Sin(x));
 XLAJIT_MAKE_UNARY(Tan, xla::Tan(x));
-REGISTER_XLA_OP(Name("Exp"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("Expm1"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("Floor"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("IsFinite"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("IsInf"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("IsNan"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Exp, xla::Exp(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Expm1, xla::Expm1(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Floor, xla::Floor(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(IsFinite, xla::IsFinite(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(IsInf, xla::IsInf(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(IsNan, xla::IsNan(x));
 // Return 1/x
 XLAJIT_MAKE_UNARY(Inv, xla::ScalarLike(x, 1.0) / x);
-REGISTER_XLA_OP(Name("Reciprocal"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(
+    Reciprocal, xla::ScalarLike(x, 1.0) / x);
 XLAJIT_MAKE_UNARY(Log, xla::Log(x));
-REGISTER_XLA_OP(Name("Log1p"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Log1p, xla::Log1p(x));
 
 XLAJIT_MAKE_UNARY(Invert, xla::Not(x));
 XLAJIT_MAKE_UNARY(LogicalNot, xla::Not(x));
@@ -85,12 +102,12 @@ XLAJIT_MAKE_UNARY(Neg, -x);
 XLAJIT_MAKE_UNARY(Rint, xla::RoundToEven(x));
 XLAJIT_MAKE_UNARY(Round, xla::RoundToEven(x));
 
-REGISTER_XLA_OP(Name("Rsqrt"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Rsqrt, xla::Rsqrt(x));
 
-REGISTER_XLA_OP(Name("Sigmoid"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Sigmoid, xla::Logistic(x));
 
 // Returns NaN if x is NaN, 0 if x is 0, -1 if x < 0 and 1 if x > 0.
-REGISTER_XLA_OP(Name("Sign"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Sign, xla::Sign(x));
 XLAJIT_MAKE_UNARY(Sinh, xla::Sinh(x));
 
 static xla::XlaOp Softplus(xla::XlaBuilder* b, xla::XlaOp features) {
@@ -115,11 +132,11 @@ XLAJIT_MAKE_UNARY(Softplus, Softplus(b, x));
 
 // softsign(x) = x / (abs(x) + 1)
 XLAJIT_MAKE_UNARY(Softsign, x / (xla::Abs(x) + xla::ScalarLike(x, 1.0)));
-REGISTER_XLA_OP(Name("Sqrt"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Sqrt, xla::Sqrt(x));
 XLAJIT_MAKE_UNARY(Square, x* x);
-REGISTER_XLA_OP(Name("Tanh"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("Real"), MlirXlaOpKernel);
-REGISTER_XLA_OP(Name("Imag"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Tanh, xla::Tanh(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Real, xla::Real(x));
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Imag, xla::Imag(x));
 XLAJIT_MAKE_UNARY(Erf, xla::Erf(x));
 XLAJIT_MAKE_UNARY(Erfc, xla::Erfc(x));
 XLAJIT_MAKE_UNARY(Erfinv, xla::ErfInv(x));
@@ -127,7 +144,7 @@ XLAJIT_MAKE_UNARY(Erfinv, xla::ErfInv(x));
 XLAJIT_MAKE_UNARY(Ndtri, xla::ScalarLike(x, std::sqrt(2.0)) *
                              xla::ErfInv(xla::ScalarLike(x, 2.0) * x -
                                          xla::ScalarLike(x, 1.0)));
-REGISTER_XLA_OP(Name("Lgamma"), MlirXlaOpKernel);
+XLAJIT_MAKE_UNARY_WITH_MLIR_FALLBACK(Lgamma, xla::Lgamma(x));
 XLAJIT_MAKE_UNARY(Digamma, xla::Digamma(x));
 XLAJIT_MAKE_UNARY(BesselI0e, xla::BesselI0e(x));
 XLAJIT_MAKE_UNARY(BesselI1e, xla::BesselI1e(x));
diff --git a/tensorflow/compiler/tf2xla/kernels/unique_op.cc b/tensorflow/compiler/tf2xla/kernels/unique_op.cc
index 46de3dd89b6115..6590dd2f55cf78 100644
--- a/tensorflow/compiler/tf2xla/kernels/unique_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/unique_op.cc
@@ -83,9 +83,12 @@ class UniqueOpBase : public XlaOpKernel {
   //
   // This is implemented as an hlo while loop.
   xla::XlaOp RollingSelectR1(XlaOpKernelContext* ctx, xla::XlaOp data,
-                             xla::XlaOp mask, int64_t size) {
+                             xla::XlaOp mask, int64_t size,
+                             const xla::DExpr& expr) {
     xla::XlaComputation cond, body;
-    const xla::Shape r1_shape = xla::ShapeUtil::MakeShape(xla::S32, {size});
+    xla::Shape r1_shape = xla::ShapeUtil::MakeShape(xla::S32, {size});
+    r1_shape.set_expression(0, expr);
+
     const xla::Shape counter_shape = xla::ShapeUtil::MakeScalarShape(xla::S32);
     const xla::Shape& single_element_shape = counter_shape;
 
@@ -136,7 +139,7 @@ class UniqueOpBase : public XlaOpKernel {
     }
 
     auto zero = xla::Zero(ctx->builder(), xla::S32);
-    auto zero_broadcast = xla::Broadcast(zero, {size});
+    auto zero_broadcast = xla::Broadcast(zero, {size}, {expr});
     auto init = xla::Tuple(ctx->builder(), {zero, data, mask, zero_broadcast});
     return xla::GetTupleElement(xla::While(cond, body, init), 3);
   }
@@ -153,13 +156,19 @@ class UniqueOpBase : public XlaOpKernel {
     auto aux = MoveAxis(input, axis, 0, input_shape);
     auto aux_shape = ctx->builder()->GetShape(aux).value();
     int64_t leading_size = aux_shape.dimensions(0);
+    auto leading_expr = aux_shape.expressions(0);
     int64_t product = 1;
+    auto product_expr = xla::DExpr::Const(1);
     for (int64_t i = 1; i < aux_shape.dimensions().size(); ++i) {
       product *= aux_shape.dimensions(i);
+      product_expr = product_expr * aux_shape.expressions(i);
     }
-    aux = xla::Reshape(aux, {leading_size, product});
+    product_expr = product_expr.simplify();
+    aux = xla::Reshape(aux, {leading_size, product},
+                       {leading_expr, product_expr});
     if (leading_size == 0) {
-      auto result_data = xla::Reshape(aux, aux_shape.dimensions());
+      auto result_data =
+          xla::Reshape(aux, aux_shape.dimensions(), aux_shape.expressions());
       result_data = MoveAxis(result_data, 0, axis, aux_shape);
       ctx->SetOutput(0, result_data);
       ctx->SetOutput(1, xla::Iota(ctx->builder(), xla::S32, leading_size));
@@ -169,12 +178,18 @@ class UniqueOpBase : public XlaOpKernel {
     sort_keys.reserve(product + 1);
     std::vector<xla::PrimitiveType> sort_types;
     sort_types.reserve(product + 1);
+    xla::Shape leading_shape = xla::ShapeUtil::MakeShape(
+        input_shape.element_type(), {leading_size},
+        std::vector<xla::DExpr>{leading_expr});
     for (int64_t i = 0; i < product; ++i) {
       xla::XlaOp slice = xla::SliceInDim(aux, i, i + 1, 1, 1);
-      sort_keys.push_back(xla::Reshape(slice, {leading_size}));
+      sort_keys.push_back(xla::Reshape(leading_shape, slice));
       sort_types.push_back(input_shape.element_type());
     }
-    auto iota = xla::Iota(ctx->builder(), xla::S32, leading_size);
+    xla::Shape iota_shape = xla::ShapeUtil::MakeShape(
+        xla::S32, {leading_size}, std::vector<xla::DExpr>{leading_expr});
+    iota_shape.set_expression(0, leading_expr);
+    auto iota = xla::Iota(ctx->builder(), iota_shape, 0);
     sort_keys.push_back(iota);
     sort_types.push_back(xla::S32);
 
@@ -202,16 +217,20 @@ class UniqueOpBase : public XlaOpKernel {
     gather_dim_numbers.add_collapsed_slice_dims(0);
     auto permuted = xla::Gather(aux, perm, gather_dim_numbers, {1, product});
     // Tail is everything except for first element.
-    auto tail = xla::SliceInDim(permuted, 1, leading_size, 1, 0);
+    auto tail = xla::SliceInDim(permuted, 1, leading_size,
+                                xla::DExpr::Const(1), leading_expr, 1, 0);
     // Init is everything except for last element.
-    auto init = xla::SliceInDim(permuted, 0, leading_size - 1, 1, 0);
+    auto init = xla::SliceInDim(permuted, 0, leading_size - 1,
+                                xla::DExpr::Const(0),
+                                (leading_expr - xla::DExpr::Const(1)).simplify(),
+                                1, 0);
     auto ne = xla::Compare(tail, init, xla::ComparisonDirection::kNe);
     auto reduce =
         xla::Reduce(ne, xla::ConstantR0(ctx->builder(), false),
                     CreateScalarOrComputation(xla::PRED, ctx->builder()), {1});
     auto mask = xla::ConvertElementType(reduce, xla::S32);
     mask = xla::PadInDim(mask, xla::One(ctx->builder(), xla::S32), 0, 1, 0);
-    auto iperm = RollingSelectR1(ctx, perm, mask, leading_size);
+    auto iperm = RollingSelectR1(ctx, perm, mask, leading_size, leading_expr);
 
     auto sort_by_iperm =
         xla::Sort({iperm, mask, perm},
@@ -232,12 +251,14 @@ class UniqueOpBase : public XlaOpKernel {
         /*is_stable=*/true);
     auto mask_permute = xla::GetTupleElement(mask_sort, 1);
     permuted = xla::Gather(aux, mask_permute, gather_dim_numbers, {1, product});
-    auto result_data = xla::Reshape(permuted, aux_shape.dimensions());
+    auto result_data = xla::Reshape(aux_shape, permuted);
     result_data = MoveAxis(result_data, 0, axis, aux_shape);
     result_data = xla::SetDimensionSize(result_data, dynamic_size, axis);
     ctx->SetOutput(0, result_data);
     auto imask = CumSumR1(ctx, mask, leading_size);
-    imask = xla::Sub(imask, xla::One(ctx->builder(), xla::S32), {});
+    auto one = xla::One(ctx->builder(), xla::S32);
+    auto one_broadcast = xla::Broadcast(one, {leading_size}, {leading_expr});
+    imask = xla::Sub(imask, one_broadcast, {});
     auto idx = xla::GetTupleElement(
         xla::Sort({perm_sort, imask},
                   xla::CreateScalarLtComputation({xla::S32, xla::S32},
diff --git a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc
index cca29f7f585907..55899c7f7b7d95 100644
--- a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc
@@ -60,16 +60,24 @@ class UnpackOp : public XlaOpKernel {
     std::vector<int64_t> start_indices(input_shape.dims(), 0);
     std::vector<int64_t> limit_indices(input_shape.dims());
     std::vector<int64_t> strides(input_shape.dims(), 1);
+    std::vector<xla::DExpr> start_exprs(input_shape.dims(), xla::DExpr::Const(0));
+    std::vector<xla::DExpr> limit_exprs;
+    limit_exprs.reserve(input_shape.dims());
     for (int i = 0; i < input_shape.dims(); ++i) {
       limit_indices[i] = input_shape.dim_size(i);
+      limit_exprs.push_back(input_shape.get_filled_expression(i));
     }
 
     for (int i = 0; i < num; ++i) {
       start_indices[axis] = i;
       limit_indices[axis] = i + 1;
-      auto slice = xla::Slice(input, start_indices, limit_indices, strides);
+      start_exprs[axis] = xla::DExpr::Const(i);
+      limit_exprs[axis] = xla::DExpr::Const(i + 1);
+      auto slice = xla::Slice(input, start_indices, limit_indices, start_exprs,
+                              limit_exprs, strides);
       // Reshape to drop the 'axis' dimension.
-      auto result = xla::Reshape(slice, output_shape.dim_sizes());
+      auto result = xla::Reshape(slice, output_shape.dim_sizes(),
+                                 output_shape.get_filled_expressions());
       ctx->SetOutput(i, result);
     }
   }
diff --git a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc
index a7a1a438f95b9e..b588525f3831a0 100644
--- a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h"
 #include "tensorflow/compiler/tf2xla/kernels/shape_util.h"
 #include "tensorflow/compiler/tf2xla/lib/scatter.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "tensorflow/compiler/tf2xla/xla_resource.h"
@@ -50,6 +51,19 @@ absl::Status ValidateAssignUpdateVariableOpShapes(XlaOpKernelContext* ctx) {
   return absl::OkStatus();
 }
 
+std::vector<xla::DExpr> BuildVariableShapeContents(const TensorShape& shape) {
+  std::vector<xla::DExpr> contents;
+  contents.reserve(shape.dims());
+  for (int i = 0; i < shape.dims(); ++i) {
+    xla::DExpr expr = shape.get_filled_expression(i);
+    contents.push_back(expr && expr->is_dynamic()
+                           ? expr
+                           : xla::DExpr::Unknown(
+                                 xla::kUnknownContentSentinel));
+  }
+  return contents;
+}
+
 class VarIsInitializedOp : public XlaOpKernel {
  public:
   explicit VarIsInitializedOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
@@ -75,7 +89,11 @@ class VariableShapeOp : public XlaOpKernel {
                    ctx->GetVariableTypeAndShape(0, &variable_dtype, &shape));
     Tensor shape_constant(out_dtype_, TensorShape({shape.dims()}));
     OP_REQUIRES_OK(ctx, TensorShapeToConstant(shape, &shape_constant));
-    ctx->SetConstantOutput(0, shape_constant);
+    auto output = XlaExpression::Constant(shape_constant);
+    if (SymbolicContentEnabled()) {
+      output.set_contents(BuildVariableShapeContents(shape));
+    }
+    ctx->SetOutputExpression(0, output);
   }
 
  private:
diff --git a/tensorflow/compiler/tf2xla/kernels/where_op.cc b/tensorflow/compiler/tf2xla/kernels/where_op.cc
index f97e6d5077efa7..4920a106808605 100644
--- a/tensorflow/compiler/tf2xla/kernels/where_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/where_op.cc
@@ -162,10 +162,16 @@ absl::StatusOr<XlaOp> CompileWhereWithSort(XlaOpKernelContext* ctx) {
   TF_ASSIGN_OR_RETURN(xla::Shape input_shape,
                       ctx->builder()->GetShape(condition));
   auto iota_shape =
-      xla::ShapeUtil::MakeShape(xla::S32, input_shape.dimensions());
+      xla::ShapeUtil::MakeShape(xla::S32, input_shape.dimensions(),
+                                input_shape.expressions());
 
   int64_t flattened_size = xla::Product(iota_shape.dimensions());
-  XlaOp reshaped_condition = xla::Reshape(condition, {flattened_size});
+  xla::DExpr flattened_expr = xla::DExpr::Const(1);
+  for (auto e : iota_shape.expressions()){
+    flattened_expr = flattened_expr * e;
+  }
+  XlaOp reshaped_condition =
+      xla::Reshape(condition, {flattened_size}, {flattened_expr});
   XlaOp zeros = xla::ZerosLike(reshaped_condition);
   XlaOp compared = xla::Ne(reshaped_condition, zeros);
 
@@ -175,7 +181,7 @@ absl::StatusOr<XlaOp> CompileWhereWithSort(XlaOpKernelContext* ctx) {
   // indices of each element.
   for (int64_t axis = 0; axis < iota_shape.dimensions_size(); ++axis) {
     XlaOp iota = xla::Iota(ctx->builder(), iota_shape, axis);
-    XlaOp reshaped = xla::Reshape(iota, {flattened_size});
+    XlaOp reshaped = xla::Reshape(iota, {flattened_size}, {flattened_expr});
     to_sort.push_back(reshaped);
     types_to_sort.push_back(xla::S32);
   }
@@ -186,7 +192,9 @@ absl::StatusOr<XlaOp> CompileWhereWithSort(XlaOpKernelContext* ctx) {
   std::vector<XlaOp> to_concat;
   for (int64_t i = 0; i < iota_shape.dimensions_size(); ++i) {
     XlaOp index_single_dim = xla::GetTupleElement(sorted, i + 1);
-    to_concat.push_back(xla::Reshape(index_single_dim, {flattened_size, 1}));
+    to_concat.push_back(xla::Reshape(index_single_dim, {flattened_size, 1},
+                                     {flattened_expr,
+                                      xla::DExpr::Const(1)}));
   }
 
   XlaOp result = xla::ConcatInDim(ctx->builder(), to_concat, 1);
@@ -214,7 +222,12 @@ absl::StatusOr<XlaOp> CompileWhereWithPrefixSum(XlaOpKernelContext* ctx) {
   TF_ASSIGN_OR_RETURN(xla::Shape input_shape, b->GetShape(condition));
 
   int64_t flattened_size = xla::Product(input_shape.dimensions());
-  XlaOp reshaped_condition = xla::Reshape(condition, {flattened_size});
+  xla::DExpr flattened_expr = xla::DExpr::Const(1);
+  for (auto e : input_shape.expressions()) {
+    flattened_expr = flattened_expr * e;
+  }
+  XlaOp reshaped_condition =
+      xla::Reshape(condition, {flattened_size}, {flattened_expr});
   XlaOp zeros = xla::ZerosLike(reshaped_condition);
   XlaOp preds =
       xla::ConvertElementType(xla::Ne(reshaped_condition, zeros), S32);
@@ -276,11 +289,16 @@ absl::StatusOr<XlaOp> CompileWhereWithPrefixSum(XlaOpKernelContext* ctx) {
   //
   // and then scatter iotas[out_idxs] into the output.
   std::vector<XlaOp> iotas_to_concat;
-  auto iota_shape = xla::ShapeUtil::MakeShape(S32, input_shape.dimensions());
+  auto iota_shape = xla::ShapeUtil::MakeShape(
+      S32, input_shape.dimensions(), input_shape.expressions());
   iotas_to_concat.reserve(iota_shape.dimensions_size());
   for (int64_t axis = 0; axis < iota_shape.dimensions_size(); ++axis) {
-    iotas_to_concat.push_back(
-        xla::Reshape(xla::Iota(b, iota_shape, axis), {flattened_size, 1}));
+    XlaOp flattened_iota =
+        xla::Reshape(xla::Iota(b, iota_shape, axis), {flattened_size},
+                     {flattened_expr});
+    iotas_to_concat.push_back(xla::Reshape(
+        flattened_iota, {flattened_size, 1},
+        {flattened_expr, xla::DExpr::Const(1)}));
   }
   XlaOp iotas = xla::ConcatInDim(b, iotas_to_concat, /*dimension=*/1);
 
@@ -305,7 +323,12 @@ absl::StatusOr<XlaOp> CompileWhereWithPrefixSum(XlaOpKernelContext* ctx) {
   XlaOp scattered = xla::Scatter(
       /*input=*/xla::Zeros(
           b, /*shape=*/xla::ShapeUtil::MakeShape(
-              S32, {flattened_size, iota_shape.dimensions_size()})),
+              S32,
+              std::vector<int64_t>{flattened_size,
+                                   iota_shape.dimensions_size()},
+              std::vector<xla::DExpr>{
+                  flattened_expr,
+                  xla::DExpr::Const(iota_shape.dimensions_size())})),
       /*scatter_indices=*/out_idxs, /*updates=*/iotas,
       /*update_computation=*/assn_computation, scatter_dnums,
       /*indices_are_sorted=*/true, /*unique_indices=*/true);
diff --git a/tensorflow/compiler/tf2xla/layout_util.cc b/tensorflow/compiler/tf2xla/layout_util.cc
index b000c49f1f962e..1db86361065df9 100644
--- a/tensorflow/compiler/tf2xla/layout_util.cc
+++ b/tensorflow/compiler/tf2xla/layout_util.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/tf2xla/shape_util.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
 #include "tensorflow/compiler/tf2xla/xla_argument.h"
@@ -131,8 +132,12 @@ absl::StatusOr<xla::XlaOp> ReshapeWithCorrectRepresentationAndSharding(
         hlo_sharding, fast_mem, shape_determination_fns, &to_shape));
   }
   if (xla::ShapeUtil::Compatible(original_shape, to_shape)) {
+    MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
     for (int64_t i = 0; i < original_shape.dimensions().size(); ++i) {
       to_shape.set_dynamic_dimension(i, original_shape.is_dynamic_dimension(i));
+      if (flags->tf_xla_enable_dynamic_sizes) {
+        to_shape.set_expression(i, original_shape.expressions(i));
+      }
     }
   }
   return xla::Reshape(to_shape, original);
diff --git a/tensorflow/compiler/tf2xla/lib/broadcast.cc b/tensorflow/compiler/tf2xla/lib/broadcast.cc
index f815b91d04be33..07d09559b15d35 100644
--- a/tensorflow/compiler/tf2xla/lib/broadcast.cc
+++ b/tensorflow/compiler/tf2xla/lib/broadcast.cc
@@ -32,9 +32,10 @@ limitations under the License.
 
 namespace tensorflow {
 
-absl::StatusOr<xla::XlaOp> BroadcastTo(xla::XlaOp input,
-                                       absl::Span<int64_t const> output_dims) {
-  return xla::BroadcastTo(input, output_dims);
+absl::StatusOr<xla::XlaOp> BroadcastTo(
+    xla::XlaOp input, absl::Span<int64_t const> output_dims,
+    absl::Span<const xla::DExpr> output_exprs) {
+  return xla::BroadcastTo(input, output_dims, output_exprs);
 }
 
 absl::Status BroadcastOpsToSame(xla::XlaOp* lhs, xla::XlaOp* rhs) {
diff --git a/tensorflow/compiler/tf2xla/lib/broadcast.h b/tensorflow/compiler/tf2xla/lib/broadcast.h
index 60630971e27466..b5903775679b30 100644
--- a/tensorflow/compiler/tf2xla/lib/broadcast.h
+++ b/tensorflow/compiler/tf2xla/lib/broadcast.h
@@ -29,8 +29,9 @@ namespace tensorflow {
 
 // Forwards to xla::BroadcastTo.
 // TODO(cheshire): Call the underlying function directly.
-absl::StatusOr<xla::XlaOp> BroadcastTo(xla::XlaOp input,
-                                       absl::Span<int64_t const> output_dims);
+absl::StatusOr<xla::XlaOp> BroadcastTo(
+    xla::XlaOp input, absl::Span<int64_t const> output_dims,
+    absl::Span<const xla::DExpr> output_exprs = {});
 
 // Forwards to xla::BroadcastOpsToSame.
 absl::Status BroadcastOpsToSame(xla::XlaOp* lhs, xla::XlaOp* rhs);
diff --git a/tensorflow/compiler/tf2xla/lib/data_format.cc b/tensorflow/compiler/tf2xla/lib/data_format.cc
index 2473b97af4c2bd..b59dbf59e54120 100644
--- a/tensorflow/compiler/tf2xla/lib/data_format.cc
+++ b/tensorflow/compiler/tf2xla/lib/data_format.cc
@@ -53,7 +53,13 @@ absl::StatusOr<xla::XlaOp> Contract(xla::XlaOp input, int64_t dim) {
                                         input_shape.dimensions().end() - 1);
   contracted_shape[dim] *= 4;
 
-  return xla::Reshape(xla::Transpose(input, permutation), contracted_shape);
+  std::vector<xla::DExpr> contracted_exprs(
+      input_shape.expressions().begin(), input_shape.expressions().end() - 1);
+  contracted_exprs[dim] =
+      contracted_exprs[dim] * xla::DExpr::Const(4);
+
+  return xla::Reshape(xla::Transpose(input, permutation), contracted_shape,
+                      contracted_exprs);
 }
 
 absl::StatusOr<xla::XlaOp> Expand(xla::XlaOp input, int64_t dim) {
@@ -85,7 +91,7 @@ absl::StatusOr<xla::XlaOp> Expand(xla::XlaOp input, int64_t dim) {
   }
   permutation.push_back(dim + 1);
 
-  return xla::Transpose(xla::Reshape(input, expanded_shape), permutation);
+  return xla::Transpose(xla::Reshape(input, expanded_shape, {}), permutation);
 }
 
 }  // namespace
diff --git a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc
index b1a93508d92896..28d3447918a2c7 100644
--- a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc
+++ b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
 
 #include <string>
+#include <variant>
 
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
@@ -35,7 +36,9 @@ limitations under the License.
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/resource_base.h"
 #include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/refcount.h"
 #include "tensorflow/core/platform/status.h"
@@ -68,6 +71,77 @@ class MLIRContextResource : public ResourceBase {
   mlir::MLIRContext mlir_ctx_;
 };
 
+bool HasDynamicExpressions(const TensorShape& shape) {
+  for (const auto& expr : shape.get_expressions()) {
+    if (expr && expr->is_dynamic()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool HasDynamicExpressions(const PartialTensorShape& shape) {
+  for (const auto& expr : shape.get_expressions()) {
+    if (expr && expr->is_dynamic()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool HasDynamicExpressions(const xla::Shape& shape) {
+  for (const auto& expr : shape.expressions()) {
+    if (expr && expr->is_dynamic()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+absl::Status RejectDynamicShapeExpressionsInMlirXlaOpKernel(
+    llvm::ArrayRef<XlaCompiler::Argument> args, const Graph& graph) {
+  for (int i = 0; i < args.size(); ++i) {
+    const auto& shape = args[i].shape;
+    const bool has_dynamic_exprs =
+        std::holds_alternative<TensorShape>(shape)
+            ? HasDynamicExpressions(std::get<TensorShape>(shape))
+            : HasDynamicExpressions(std::get<xla::Shape>(shape));
+    if (has_dynamic_exprs) {
+      return errors::Unimplemented(
+          "MlirXlaOpKernel does not support dynamic shape expressions. "
+          "Argument ",
+          i, " carries a dynamic expression.");
+    }
+  }
+
+  for (Node* node : graph.nodes()) {
+    for (const auto& name_attr_pair : node->attrs()) {
+      const auto& attr_name = name_attr_pair.first;
+      const auto& attr_value = name_attr_pair.second;
+      auto maybe_reject_shape = [&](const TensorShapeProto& shape_proto) {
+        const PartialTensorShape shape(shape_proto);
+        if (!HasDynamicExpressions(shape)) {
+          return absl::OkStatus();
+        }
+        return errors::Unimplemented(
+            "MlirXlaOpKernel does not support dynamic shape expressions. "
+            "Node '",
+            node->name(), "' attribute '", attr_name,
+            "' carries a dynamic expression.");
+      };
+
+      if (attr_value.value_case() == AttrValue::kShape) {
+        TF_RETURN_IF_ERROR(maybe_reject_shape(attr_value.shape()));
+      } else if (attr_value.value_case() == AttrValue::kList) {
+        for (const auto& shape_proto : attr_value.list().shape()) {
+          TF_RETURN_IF_ERROR(maybe_reject_shape(shape_proto));
+        }
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 absl::Status MlirXlaOpKernel::ContextToXlaArgs(
@@ -143,6 +217,8 @@ absl::Status MlirXlaOpKernel::ConstructXlaOp(XlaOpKernelContext* ctx) {
   // Create a graph that wraps the kernel.
   TF_ASSIGN_OR_RETURN(auto graph,
                       CreateSingleOpGraph(def(), xla_args, result_dtypes));
+  TF_RETURN_IF_ERROR(
+      RejectDynamicShapeExpressionsInMlirXlaOpKernel(xla_args, *graph));
 
   ResourceMgr* res_manager = ctx->op_kernel_context()->resource_manager();
   MLIRContextResource* ctx_res;
diff --git a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h
index 6053f5d68635d0..7b70096851b3df 100644
--- a/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h
+++ b/tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2XLA_MLIR_XLA_OP_KERNEL_H_
 #define TENSORFLOW_COMPILER_TF2XLA_MLIR_XLA_OP_KERNEL_H_
 
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -36,6 +37,14 @@ class MlirXlaOpKernel : public XlaOpKernel {
   absl::Status ConstructXlaOp(XlaOpKernelContext* ctx);
 };
 
+template <typename NativeOp>
+OpKernel* CreateDynamicNativeXlaOpKernel(OpKernelConstruction* ctx) {
+  if (GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes) {
+    return new NativeOp(ctx);
+  }
+  return new MlirXlaOpKernel(ctx);
+}
+
 }  // namespace tensorflow
 
 #endif  // TENSORFLOW_COMPILER_TF2XLA_MLIR_XLA_OP_KERNEL_H_
diff --git a/tensorflow/compiler/tf2xla/ops/BUILD b/tensorflow/compiler/tf2xla/ops/BUILD
index 61a329aa46f1ba..6014b73c35949c 100644
--- a/tensorflow/compiler/tf2xla/ops/BUILD
+++ b/tensorflow/compiler/tf2xla/ops/BUILD
@@ -16,6 +16,7 @@ cc_library(
     name = "xla_ops",
     srcs = ["xla_ops.cc"],
     deps = [
+        "//tensorflow/compiler/jit:flags",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
diff --git a/tensorflow/compiler/tf2xla/ops/xla_ops.cc b/tensorflow/compiler/tf2xla/ops/xla_ops.cc
index 6a67cfa237af70..516ae853d40290 100644
--- a/tensorflow/compiler/tf2xla/ops/xla_ops.cc
+++ b/tensorflow/compiler/tf2xla/ops/xla_ops.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_join.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "xla/service/shape_inference.h"
 #include "xla/shape.h"
 #include "xla/xla_data.pb.h"
@@ -1143,16 +1144,32 @@ xla::Shape GetShape(shape_inference::ShapeHandle shape_handle,
   }
   std::vector<int64_t> dims;
   std::vector<bool> dynamic_dims;
+  std::vector<xla::DExpr> expressions;
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
   for (int i = 0, rank = c->Rank(shape_handle); i < rank; ++i) {
-    bool is_dynamic = !c->ValueKnown(c->Dim(shape_handle, i));
+    const auto dim = c->Dim(shape_handle, i);
+    const bool is_dynamic = !c->ValueKnown(dim);
     dynamic_dims.push_back(is_dynamic);
-    dims.push_back(is_dynamic ? xla::Shape::kUnboundedSize
-                              : c->Value(c->Dim(shape_handle, i)));
+    if (flags->tf_xla_enable_dynamic_sizes) {
+      DimExpr* expr = c->GetDimExpr(dim);
+      expressions.push_back(
+          expr != nullptr
+              ? *expr
+              : c->ValueKnown(dim)
+                    ? xla::DExpr::Const(c->Value(dim))
+                    : xla::DExpr::Unknown(xla::kMissingExpressionSentinel));
+    }
+    dims.push_back(is_dynamic ? xla::Shape::kUnboundedSize : c->Value(dim));
   }
-  return xla::Shape(
+  xla::Shape sh(
       // Type matters only for indices. S64 is the widest possible type.
       xla::PrimitiveType::S64, dims,
       absl::InlinedVector<bool, 4>(dynamic_dims.begin(), dynamic_dims.end()));
+
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    sh.set_expressions(expressions);
+  }
+  return sh;
 }
 
 REGISTER_OP("XlaGather")
diff --git a/tensorflow/compiler/tf2xla/shape_util.cc b/tensorflow/compiler/tf2xla/shape_util.cc
index 0d7549d81c20f6..132daea9ee73e7 100644
--- a/tensorflow/compiler/tf2xla/shape_util.cc
+++ b/tensorflow/compiler/tf2xla/shape_util.cc
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/tf2xla/type_util.h"
 #include "xla/layout_util.h"
 #include "xla/shape_util.h"
@@ -100,6 +101,11 @@ absl::Status XLAShapeToTensorShape(const xla::Shape& shape,
   for (int i = 0; i < shape.dimensions().size(); ++i) {
     TF_RETURN_IF_ERROR(tensor_shape->AddDimWithStatus(shape.dimensions(i)));
   }
+  if (!shape.expressions().empty()) {
+    std::vector<xla::DExpr> dexprs(shape.expressions().begin(),
+                                   shape.expressions().end());
+    tensor_shape->set_expressions(std::move(dexprs));
+  }
   return absl::OkStatus();
 }
 
@@ -168,6 +174,11 @@ xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type,
   int rank = tensor_shape.dims();
   std::vector<int64_t> dimensions(rank);
   std::vector<int64_t> layout(rank);
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
+  std::vector<xla::DExpr> expressions;
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    expressions.resize(rank);
+  }
   for (int d = 0; d < rank; ++d) {
     dimensions[d] = tensor_shape.dim_size(d);
     if (dimensions[d] < 0) {
@@ -175,11 +186,17 @@ xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type,
                       "shape; returning unknown sentinel value";
       return xla::ShapeUtil::MakeShapeWithDenseLayout(type, {0}, {0});
     }
+    if (flags->tf_xla_enable_dynamic_sizes) {
+      expressions[d] = tensor_shape.get_filled_expression(d);
+    }
   }
   // XLA uses minor-to-major; Tensorflow uses major-to-minor.
   std::iota(layout.rbegin(), layout.rend(), 0);
   xla::Shape result =
       xla::ShapeUtil::MakeShapeWithDenseLayout(type, dimensions, layout);
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    result.set_expressions(expressions);
+  }
   return result;
 }
 
@@ -200,18 +217,35 @@ absl::StatusOr<xla::Shape> TensorShapeToXLAShape(
   return out;
 }
 
+inline static int var_id = 1;
+
 xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type,
                                  const TensorShape& tensor_shape) {
   int rank = tensor_shape.dims();
   std::vector<int64_t> dimensions(rank);
   std::vector<int64_t> layout(rank);
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
+  std::vector<xla::DExpr> expressions;
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    expressions.resize(rank);
+  }
+
   for (int d = 0; d < rank; ++d) {
     dimensions[d] = tensor_shape.dim_size(d);
+    if (flags->tf_xla_enable_dynamic_sizes) {
+      expressions[d] = tensor_shape.get_filled_expression(d);
+    }
   }
+
   // XLA uses minor-to-major; Tensorflow uses major-to-minor.
   std::iota(layout.rbegin(), layout.rend(), 0);
 
-  return xla::ShapeUtil::MakeShapeWithDenseLayout(type, dimensions, layout);
+  auto shape =
+      xla::ShapeUtil::MakeShapeWithDenseLayout(type, dimensions, layout);
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    shape.set_expressions(expressions);
+  }
+  return shape;
 }
 
 absl::StatusOr<std::vector<int>> GetShapeLayoutVector(const xla::Shape& shape) {
diff --git a/tensorflow/compiler/tf2xla/symbolic_content_util.h b/tensorflow/compiler/tf2xla/symbolic_content_util.h
new file mode 100644
index 00000000000000..3dd471a4e9ab02
--- /dev/null
+++ b/tensorflow/compiler/tf2xla/symbolic_content_util.h
@@ -0,0 +1,29 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_TF2XLA_SYMBOLIC_CONTENT_UTIL_H_
+#define TENSORFLOW_COMPILER_TF2XLA_SYMBOLIC_CONTENT_UTIL_H_
+
+#include "tensorflow/compiler/jit/flags.h"
+
+namespace tensorflow {
+
+inline bool SymbolicContentEnabled() {
+  return GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content;
+}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_TF2XLA_SYMBOLIC_CONTENT_UTIL_H_
diff --git a/tensorflow/compiler/tf2xla/xla_argument.cc b/tensorflow/compiler/tf2xla/xla_argument.cc
index 8b91dd3870b7d5..71ea93bf4fc628 100644
--- a/tensorflow/compiler/tf2xla/xla_argument.cc
+++ b/tensorflow/compiler/tf2xla/xla_argument.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/xla_argument.h"
 
+#include "absl/algorithm/container.h"
 #include "llvm/ADT/STLExtras.h"
 
 namespace tensorflow {
@@ -46,6 +47,14 @@ bool XlaArgument::operator==(const XlaArgument& other) const {
   if (constant_value.shape() != other.constant_value.shape()) {
     return false;
   }
+  if (!absl::c_equal(
+          constant_value_expressions, other.constant_value_expressions,
+          [](const xla::ExpressionProto& lhs,
+             const xla::ExpressionProto& rhs) {
+            return lhs.SerializeAsString() == rhs.SerializeAsString();
+          })) {
+    return false;
+  }
   if (is_same_data_across_replicas != other.is_same_data_across_replicas) {
     return false;
   }
diff --git a/tensorflow/compiler/tf2xla/xla_argument.h b/tensorflow/compiler/tf2xla/xla_argument.h
index 9e2eccd29b1885..f037cb8610619e 100644
--- a/tensorflow/compiler/tf2xla/xla_argument.h
+++ b/tensorflow/compiler/tf2xla/xla_argument.h
@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_
 #define TENSORFLOW_COMPILER_TF2XLA_XLA_ARGUMENT_H_
 
+#include <vector>
+
 #include "absl/types/optional.h"
 #include "absl/types/span.h"
 #include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h"
@@ -23,6 +25,7 @@ limitations under the License.
 #include "xla/hlo/builder/xla_builder.h"
 #include "xla/hlo/ir/hlo_sharding.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
 
 namespace tensorflow {
 
@@ -76,6 +79,10 @@ struct XlaArgument {
   // host-memory tensor.
   Tensor constant_value;
 
+  // Symbolic expressions for each element of a compile-time constant.
+  // This is only used for shape-like integer tensors crossing cluster
+  // boundaries.
+  std::vector<xla::ExpressionProto> constant_value_expressions;
   // The upper bounds of the value.
   std::optional<Tensor> value_bound;
 
@@ -116,6 +123,7 @@ struct XlaArgument {
 
   // Returns the dimension sizes for either TensorShape or xla::Shape.
   std::vector<int64_t> DimensionSizes() const;
+  std::vector<xla::DExpr> DimensionExpressions() const;
   absl::InlinedVector<int64_t, 4> DimensionSizesAsInlinedVector() const;
 
   // Returns the human-readable string for either TensorShape or xla::Shape.
diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc
index b7cff00c8a0bfe..90faea1745d327 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler.cc
@@ -32,12 +32,14 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/variant.h"
 #include "tensorflow/compiler/jit/defs.h"
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/shape_inference.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "tensorflow/compiler/jit/xla_compile_util.h"
 #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
 #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
@@ -93,7 +95,6 @@ namespace {
 constexpr char kSingleOpComponent[] = "TF2XLA_XLA_COMPILER_COMPILE_SINGLE_OP";
 constexpr char kCompileFunctionComponent[] =
     "TF2XLA_XLA_COMPILER_COMPILE_FUNCTION";
-
 // Checks that arguments `args` match types `types`.
 absl::Status CheckSignature(const DataTypeVector& types,
                             absl::Span<const XlaCompiler::Argument> args) {
@@ -511,6 +512,19 @@ std::vector<int64_t> XlaCompiler::Argument::DimensionSizes() const {
   }
 }
 
+std::vector<xla::DExpr> XlaCompiler::Argument::DimensionExpressions() const {
+  if (absl::holds_alternative<TensorShape>(shape)) {
+    return std::get<TensorShape>(shape).get_filled_expressions();
+  } else {
+    std::vector<xla::DExpr> expressions;
+    expressions.reserve(std::get<xla::Shape>(shape).expressions().size());
+    for (const auto& expr : std::get<xla::Shape>(shape).expressions()) {
+      expressions.push_back(expr);
+    }
+    return expressions;
+  }
+}
+
 absl::InlinedVector<int64_t, 4>
 XlaCompiler::Argument::DimensionSizesAsInlinedVector() const {
   if (absl::holds_alternative<TensorShape>(shape)) {
@@ -839,7 +853,9 @@ absl::Status XlaCompiler::CompileFunction(
                                      std::vector<TensorShape>{tensor_shape});
       }
     } else {
+      auto* val_ptr = std::get_if<TensorShape>(&args[i].shape);
       TensorShape tensor_shape = std::get<TensorShape>(args[i].shape);
+      AttrSlice n_attrs = fbody->arg_nodes[i]->attrs();
       fbody->arg_nodes[i]->ClearAttr("_output_shapes");
       fbody->arg_nodes[i]->AddAttr("_output_shapes",
                                    std::vector<TensorShape>{tensor_shape});
@@ -930,14 +946,17 @@ absl::Status XlaCompiler::XLAShapeForArgument(
         TF_RETURN_IF_ERROR(RewriteLayoutWithShardedShape(
             arg_sharding, /*use_fast_memory=*/false,
             options_.shape_determination_fns, xla_shape));
-        // If the arg is dynamic then we update the shape to reflect that. The
-        // layout etc above lose it by forcing a swap to TensorShape.
-        if (std::holds_alternative<xla::Shape>(arg.shape) &&
-            std::get<xla::Shape>(arg.shape).is_dynamic()) {
-          xla::Shape dynamic_shape = std::get<xla::Shape>(arg.shape);
-          for (int i = 0; i < xla_shape->dimensions().size(); ++i) {
-            xla_shape->set_dynamic_dimension(
-                i, dynamic_shape.is_dynamic_dimension(i));
+        // If the arg carries dynamic metadata or symbolic expressions then we
+        // update the shape to reflect that. The layout logic above routes
+        // through TensorShape and can otherwise discard this information.
+        if (std::holds_alternative<xla::Shape>(arg.shape)) {
+          const xla::Shape& original_shape = std::get<xla::Shape>(arg.shape);
+          if (original_shape.is_dynamic() || original_shape.has_dynamic_expr()) {
+            for (int i = 0; i < xla_shape->dimensions().size(); ++i) {
+              xla_shape->set_dynamic_dimension(
+                  i, original_shape.is_dynamic_dimension(i));
+              xla_shape->set_expression(i, original_shape.expressions(i));
+            }
           }
         }
       } else {
@@ -1083,6 +1102,7 @@ absl::Status XlaCompiler::BuildArguments(
         TF_RET_CHECK(absl::holds_alternative<TensorShape>(arg.shape));
         // TODO(phawkins): this code assumes that resource arguments do not
         // alias.
+      auto* val_ptr = std::get_if<TensorShape>(&arg.shape);
         XlaResource* resource =
             context->AddResource(std::make_unique<XlaResource>(
                 arg.resource_kind, i, arg.name, arg.type,
@@ -1108,6 +1128,25 @@ absl::Status XlaCompiler::BuildArguments(
       }
       case XlaCompiler::Argument::kConstant:
         arg_expression = XlaExpression::Constant(arg.constant_value);
+        if (!arg.constant_value_expressions.empty()) {
+          VLOG(1) << "BuildArguments attaching "
+                    << arg.constant_value_expressions.size()
+                    << " constant_value_expressions to constant arg " << i
+                    << " (" << arg.name << ")";
+          // Preserve symbolic per-element metadata for shape-like constants so
+          // later tf2xla consumers can recover dynamic contents from them.
+          std::vector<xla::DExpr> contents;
+          contents.reserve(arg.constant_value_expressions.size());
+          for (const xla::ExpressionProto& expr :
+               arg.constant_value_expressions) {
+            xla::DExpr parsed = xla::DExprFromProto(expr);
+            contents.push_back(parsed && parsed->is_dynamic()
+                                   ? std::move(parsed)
+                                   : xla::DExpr::Unknown(
+                                         xla::kUnknownContentSentinel));
+          }
+          arg_expression.set_contents(std::move(contents));
+        }
         break;
       case XlaCompiler::Argument::kInvalid:
         return errors::Internal(
@@ -1203,6 +1242,10 @@ absl::Status XlaCompiler::BuildArguments(
       xla::XlaScopedShardingAssignment assign_sharding(
           builder, it == arg_shardings.end() ? std::optional<xla::OpSharding>()
                                              : it->second);
+      auto& arg = args[input_to_args->at(i)];
+      xla::OpMetadata arg_metadata;
+      arg_metadata.set_op_name(arg.node_name);
+      builder->SetOneShotOpMetadata(arg_metadata);
       if (is_entry_computation) {
         // Add an entry to is_same_across_replicas for every leaf buffer.
         std::vector<bool> is_same_across_replicas(
@@ -1222,10 +1265,9 @@ absl::Status XlaCompiler::BuildArguments(
 
   // Fill in the handles in non-constant arguments, and reshape parameters
   // back to their correct shapes.
-  VLOG(2) << "XLA computation inputs:";
   for (std::vector<int>::size_type i = 0; i < input_to_args->size(); ++i) {
     const XlaCompiler::Argument& arg = args[input_to_args->at(i)];
-    VLOG(2) << "  XLA arg " << i
+    VLOG(2) << " XLA arg " << i
             << " shape: " << xla::ShapeUtil::HumanString(arg_shapes[i])
             << " name: " << arg.name << " TF arg " << input_to_args->at(i)
             << " node name: " << arg.node_name
@@ -1251,7 +1293,9 @@ absl::Status XlaCompiler::BuildArguments(
         // return values of functions, and then reshape unconditionally.
         if (is_entry_computation) {
           arg_expression = XlaExpression::XlaOp(
-              xla::Reshape(arg_handles[i], arg.DimensionSizes()), arg.type);
+              xla::Reshape(arg_handles[i], arg.DimensionSizes(),
+                           arg.DimensionExpressions()),
+              arg.type);
         } else {
           arg_expression = XlaExpression::XlaOp(arg_handles[i], arg.type);
           if (arg.value_bound) {
diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc
index 5aef5601af61ca..bf2728b693c9d8 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc
@@ -16,8 +16,10 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 
 #include <gmock/gmock.h>
+#include <optional>
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/cc/framework/ops.h"
 #include "tensorflow/cc/ops/const_op.h"
 #include "tensorflow/cc/ops/data_flow_ops.h"
@@ -57,8 +59,11 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/grappler/costs/graph_properties.h"
+#include "tensorflow/core/grappler/grappler_item.h"
 #include "tensorflow/core/graph/algorithm.h"
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/kernels/ops_testutil.h"
@@ -96,6 +101,42 @@ class XlaCompilerTest : public ::testing::Test {
   std::unique_ptr<FunctionLibraryDefinition> flib_def_;
 };
 
+class ScopedTfXlaDynamicSizesFlag {
+ public:
+  ScopedTfXlaDynamicSizesFlag() {
+    old_value_ = GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes;
+    old_symbolic_content_value_ =
+        GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content;
+    GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes = true;
+    GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content = true;
+    SetTensorShapeExpressionsEnabledForTesting(true);
+  }
+
+  ~ScopedTfXlaDynamicSizesFlag() {
+    SetTensorShapeExpressionsEnabledForTesting(std::nullopt);
+    GetMarkForCompilationPassFlags()->tf_xla_enable_dynamic_sizes = old_value_;
+    GetMarkForCompilationPassFlags()->tf_xla_enable_symbolic_content =
+        old_symbolic_content_value_;
+  }
+
+ private:
+  bool old_value_ = false;
+  bool old_symbolic_content_value_ = false;
+};
+
+class XlaCompilerDynamicSizesTest : public XlaCompilerTest {
+ protected:
+  void SetUp() override {
+    dynamic_sizes_flag_ = std::make_unique<ScopedTfXlaDynamicSizesFlag>();
+    XlaCompilerTest::SetUp();
+  }
+
+  void TearDown() override { dynamic_sizes_flag_.reset(); }
+
+ private:
+  std::unique_ptr<ScopedTfXlaDynamicSizesFlag> dynamic_sizes_flag_;
+};
+
 namespace {
 
 // Helper class to test the ability to pass resources through to XLA
@@ -106,193 +147,3000 @@ class DummyResourceForTest : public ResourceBase {
   void Increment() { ++value_; }
   int Get() { return value_; }
 
- private:
-  int value_ = 0;
-};
+ private:
+  int value_ = 0;
+};
+
+class DummyReadResourceOp : public XlaOpKernel {
+ public:
+  explicit DummyReadResourceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {
+    ResourceMgr* rm = ctx->op_kernel_context()->resource_manager();
+    OP_REQUIRES(ctx, rm, errors::Internal("No resource manager."));
+    DummyResourceForTest* dummy;
+    OP_REQUIRES_OK(ctx, rm->Lookup<DummyResourceForTest>(
+                            rm->default_container(), "dummy", &dummy));
+    dummy->Increment();
+    dummy->Unref();
+
+    ctx->SetOutput(0, ctx->Input(0));
+    ctx->SetOutput(1, ctx->Input(0));
+  }
+};
+
+class DummyReadResourceCC {
+ public:
+  DummyReadResourceCC(const Scope& scope, const Input& value) {
+    if (!scope.ok()) return;
+    auto _value = ops::AsNodeOut(scope, value);
+    if (!scope.ok()) return;
+    Node* ret;
+    const auto unique_name = scope.GetUniqueNameForOp("DummyReadResource");
+    auto builder = NodeBuilder(unique_name, "DummyReadResource").Input(_value);
+    scope.UpdateBuilder(&builder);
+    scope.UpdateStatus(builder.Finalize(scope.graph(), &ret));
+    if (!scope.ok()) return;
+    scope.UpdateStatus(scope.DoShapeInference(ret));
+    if (!scope.ok()) return;
+    this->output1_ = Output(ret, 0);
+    this->output2_ = Output(ret, 1);
+  }
+
+  Output output1_;
+  Output output2_;
+};
+
+REGISTER_OP("DummyReadResource")
+    .Input("input: int32")
+    .Output("output1: int32")
+    .Output("output2: int32")
+    .SetShapeFn(shape_inference::UnknownShape)
+    .Doc(R"doc(
+A dummy Op.
+
+input: dummy input.
+output1: dummy output.
+output2: dummy output.
+)doc");
+
+REGISTER_XLA_OP(Name("DummyReadResource"), DummyReadResourceOp);
+
+// DummyDuplicateOp is present purely to test multiple REGISTER_XLA_OP calls
+// on the same Op name below.
+class DummyDuplicateOp : public XlaOpKernel {
+ public:
+  explicit DummyDuplicateOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {
+    ctx->SetOutput(0, ctx->Input(0));
+  }
+};
+
+REGISTER_OP("DummyDuplicateOp")
+    .Input("input: int32")
+    .Output("output: int32")
+    .Doc(R"doc(
+A dummy Op.
+
+input: dummy input.
+output: dummy output.
+)doc");
+
+REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_CPU_XLA_JIT),
+                DummyDuplicateOp);
+REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_GPU_XLA_JIT),
+                DummyDuplicateOp);
+
+// Tests compilation and execution of an empty graph.
+TEST_F(XlaCompilerTest, EmptyReturnValues) {
+  XlaCompiler compiler(DefaultOptions());
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add",
+                                     std::move(graph),
+                                     /*args=*/{}, &result));
+
+  TF_ASSERT_OK(client_->Execute(*result.computation, {}).status());
+}
+
+// Tests compilation and execution of a graph that adds two tensors.
+TEST_F(XlaCompilerTest, Simple) {
+  // Builds a graph that adds two Tensors.
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
+  auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
+  auto c = ops::Add(scope.WithOpName("C"), a, b);
+  auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  // Builds a description of the arguments.
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = TensorShape({2});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = TensorShape({2});
+
+  // Compiles the graph.
+  XlaCompiler compiler(DefaultOptions());
+
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add",
+                                     std::move(graph), args, &result));
+
+  // Tests that the generated computation works.
+  xla::Literal param0_literal = xla::LiteralUtil::CreateR1<int32>({7, 42});
+  xla::Literal param1_literal = xla::LiteralUtil::CreateR1<int32>({-3, 101});
+  std::unique_ptr<xla::GlobalData> param0_data =
+      client_->TransferToServer(param0_literal).value();
+  std::unique_ptr<xla::GlobalData> param1_data =
+      client_->TransferToServer(param1_literal).value();
+
+  std::unique_ptr<xla::GlobalData> actual =
+      client_
+          ->Execute(*result.computation, {param0_data.get(), param1_data.get()})
+          .value();
+  xla::Literal actual_literal = client_->Transfer(*actual).value();
+
+  xla::Literal expected0 = xla::LiteralUtil::CreateR1<int32>({4, 143});
+  xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0});
+  EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal));
+}
+
+absl::StatusOr<std::unique_ptr<xla::HloModule>> LoadModuleFromHloProto(
+    const xla::HloModuleProto& module_proto) {
+  TF_ASSIGN_OR_RETURN(auto module_config,
+                      xla::HloModule::CreateModuleConfigFromProto(
+                          module_proto, xla::GetDebugOptionsFromFlags()));
+  return xla::CreateModuleFromProto(module_proto, module_config);
+}
+
+// Tests compilation and execution of a graph that adds two tensors with dynamic
+// shape parameters.
+TEST_F(XlaCompilerTest, SimpleDynamicShapeParameter) {
+  // Builds a graph that adds two Tensors.
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
+  auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
+  auto c = ops::Add(scope.WithOpName("C"), a, b);
+  auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  // Builds a description of the arguments.
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape =
+      xla::ShapeUtil::MakeShape(/*element_type=*/xla::S32, /*dimensions=*/{2},
+                                /*dynamic_dimensions=*/std::vector<bool>{true},
+                                /*expressions=*/{});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = TensorShape(/*dimensions=*/{2});
+
+  // Compiles the graph.
+  XlaCompiler compiler(DefaultOptions());
+
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add",
+                                     std::move(graph), args, &result));
+
+  auto hlo = result.computation->proto();
+  TF_ASSERT_OK_AND_ASSIGN(auto module, LoadModuleFromHloProto(hlo));
+  EXPECT_EQ(module->computation_count(), 1);
+  EXPECT_TRUE(module->mutable_computation(0)
+                  ->parameter_instruction(0)
+                  ->shape()
+                  .is_dynamic());
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, DynamicShapeParameterPreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto identity = ops::Identity(scope.WithOpName("identity"), input);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), identity, 0);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6}, std::vector<xla::DExpr>{xla::DExpr::Var(1)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "identity",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(1)));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          LoadModuleFromHloProto(result.computation->proto()));
+  const xla::Shape& param_shape =
+      module->entry_computation()->parameter_instruction(0)->shape();
+  EXPECT_TRUE(
+      xla::DynExpr::equal(param_shape.expressions(0), xla::DExpr::Var(1)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(1)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       TensorFlowAndXlaAgreeOnReshapeExpression) {
+  Scope inference_scope = Scope::NewRootScope().ExitOnError();
+  auto placeholder = ops::Placeholder(
+      inference_scope.WithOpName("placeholder"), DT_FLOAT,
+      ops::Placeholder::Shape(PartialTensorShape({-1, 24})));
+  auto inference_shape =
+      ops::Const(inference_scope.WithOpName("shape"), {-1, 12});
+  auto inferred_reshape = ops::Reshape(
+      inference_scope.WithOpName("reshape"), placeholder, inference_shape);
+
+  grappler::GrapplerItem item;
+  TF_ASSERT_OK(inference_scope.ToGraphDef(&item.graph));
+  grappler::GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(
+      /*assume_valid_feeds=*/false,
+      /*aggressive_shape_inference=*/false,
+      /*include_input_tensor_values=*/false,
+      /*include_output_tensor_values=*/false,
+      /*enable_dynamic_value_inference=*/true));
+
+  const TensorShapeProto& inferred_input_shape =
+      properties.GetOutputProperties("placeholder").at(0).shape();
+  const TensorShapeProto& inferred_output_shape =
+      properties.GetOutputProperties("reshape").at(0).shape();
+  ASSERT_EQ(inferred_input_shape.expressions_size(), 2);
+  ASSERT_EQ(inferred_output_shape.expressions_size(), 2);
+  const xla::DExpr inferred_batch_expr =
+      DimExprFromProto(inferred_input_shape.expressions(0));
+  const xla::DExpr inferred_output_expr =
+      DimExprFromProto(inferred_output_shape.expressions(0));
+  ASSERT_TRUE(inferred_batch_expr->is_dynamic());
+  EXPECT_TRUE(xla::DynExpr::equal(
+      inferred_output_expr, inferred_batch_expr * xla::DExpr::Const(2)));
+
+  Scope compilation_scope = Scope::NewRootScope().ExitOnError();
+  auto arg = ops::_Arg(compilation_scope.WithOpName("arg"), DT_FLOAT, 0);
+  auto compilation_shape =
+      ops::Const(compilation_scope.WithOpName("shape"), {-1, 12});
+  auto reshape = ops::Reshape(compilation_scope.WithOpName("reshape"), arg,
+                              compilation_shape);
+  auto retval =
+      ops::_Retval(compilation_scope.WithOpName("retval"), reshape, 0);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(compilation_scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 24},
+      std::vector<xla::DExpr>{inferred_batch_expr, xla::DExpr::Const(24)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "reshape", std::move(graph),
+                                     args, &result));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          LoadModuleFromHloProto(result.computation->proto()));
+  const xla::Shape& parameter_shape =
+      module->entry_computation()->parameter_instruction(0)->shape();
+  EXPECT_TRUE(xla::DynExpr::equal(parameter_shape.expressions(0),
+                                  inferred_batch_expr));
+  EXPECT_TRUE(xla::DynExpr::equal(parameter_shape.expressions(1),
+                                  xla::DExpr::Const(24)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_EQ(result_shape.dimensions(0), 16);
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0),
+                                  inferred_output_expr));
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1),
+                                  xla::DExpr::Const(12)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       ShapeValueChainPreservesExpressionThroughTile) {
+  // Shape turns the dynamic input dimension into the symbolic value A.
+  // Slicing, squeezing, and packing must preserve it as {A, 1}, which Tile
+  // then uses as the shape of a repeated static row.
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto arg = ops::_Arg(scope.WithOpName("arg"), DT_FLOAT, 0);
+  auto input_shape = ops::Shape(scope.WithOpName("input_shape"), arg);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0}, {1});
+  auto end = ops::Const<int32>(scope.WithOpName("end"), {1}, {1});
+  auto strides = ops::Const<int32>(scope.WithOpName("strides"), {1}, {1});
+  auto first_dim_vector = ops::StridedSlice(
+      scope.WithOpName("first_dim_vector"), input_shape, begin, end, strides);
+  auto first_dim = ops::Squeeze(scope.WithOpName("first_dim"),
+                                first_dim_vector, ops::Squeeze::Axis({0}));
+  auto one = ops::Const<int32>(scope.WithOpName("one"), 1, {});
+  auto multiples = ops::Stack(scope.WithOpName("multiples"),
+                              std::vector<Output>{first_dim, one});
+  auto base_row = ops::Const<float>(
+      scope.WithOpName("base_row"),
+      {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+       1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+      {1, 24});
+  auto tile = ops::Tile(scope.WithOpName("tile"), base_row, multiples);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), tile, 0);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  // Compile with a physical bound of 8 while keeping its logical size as A.
+  const xla::DExpr input_expr = xla::DExpr::Var(1);
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 24},
+      std::vector<xla::DExpr>{input_expr, xla::DExpr::Const(24)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "tile",
+                                     std::move(graph), args, &result));
+
+  // The physical output remains bounded by 8, but its first dimension must
+  // still describe the runtime value A rather than the bound.
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_EQ(result_shape.dimensions(0), 8);
+  EXPECT_EQ(result_shape.dimensions(1), 24);
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), input_expr));
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1),
+                                  xla::DExpr::Const(24)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       SizeValueChainPreservesExpressionThroughRange) {
+  // Size returns a scalar, so the relationship to A is carried as symbolic
+  // value content rather than as a dimension expression. Range must consume
+  // that content when it infers the length of its one-dimensional output.
+  Scope inference_scope = Scope::NewRootScope().ExitOnError();
+  auto placeholder = ops::Placeholder(
+      inference_scope.WithOpName("placeholder"), DT_FLOAT,
+      ops::Placeholder::Shape(PartialTensorShape({-1, 1})));
+  auto size = ops::Size(inference_scope.WithOpName("size"), placeholder);
+  auto zero = ops::Const<int32>(inference_scope.WithOpName("zero"), 0, {});
+  auto three = ops::Const<int32>(inference_scope.WithOpName("three"), 3, {});
+  auto inferred_range =
+      ops::Range(inference_scope.WithOpName("range"), zero, size, three);
+
+  grappler::GrapplerItem item;
+  TF_ASSERT_OK(inference_scope.ToGraphDef(&item.graph));
+  grappler::GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(
+      /*assume_valid_feeds=*/false,
+      /*aggressive_shape_inference=*/false,
+      /*include_input_tensor_values=*/false,
+      /*include_output_tensor_values=*/false,
+      /*enable_dynamic_value_inference=*/true));
+
+  const TensorShapeProto& inferred_input_shape =
+      properties.GetOutputProperties("placeholder").at(0).shape();
+  const TensorShapeProto& inferred_range_shape =
+      properties.GetOutputProperties("range").at(0).shape();
+  ASSERT_EQ(inferred_input_shape.expressions_size(), 2);
+  ASSERT_EQ(inferred_range_shape.expressions_size(), 1);
+  const xla::DExpr input_expr =
+      DimExprFromProto(inferred_input_shape.expressions(0));
+  const xla::DExpr expected_output_expr =
+      DimExprFromProto(inferred_range_shape.expressions(0));
+  // With start 0 and delta 3, the number of elements is ceil(A / 3).
+  EXPECT_TRUE(xla::DynExpr::equal(
+      expected_output_expr,
+      ((input_expr + xla::DExpr::Const(2)) / xla::DExpr::Const(3))
+          .simplify()));
+
+  // Compile the equivalent graph and check the same expression reaches HLO.
+  Scope compilation_scope = Scope::NewRootScope().ExitOnError();
+  auto arg = ops::_Arg(compilation_scope.WithOpName("arg"), DT_FLOAT, 0);
+  auto compiled_size = ops::Size(compilation_scope.WithOpName("size"), arg);
+  auto compiled_zero =
+      ops::Const<int32>(compilation_scope.WithOpName("zero"), 0, {});
+  auto compiled_three =
+      ops::Const<int32>(compilation_scope.WithOpName("three"), 3, {});
+  auto range = ops::Range(compilation_scope.WithOpName("range"), compiled_zero,
+                          compiled_size, compiled_three);
+  auto retval =
+      ops::_Retval(compilation_scope.WithOpName("retval"), range, 0);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(compilation_scope.ToGraph(graph.get()));
+
+  // Reuse the variable inferred by TensorFlow at the tf2xla boundary so both
+  // layers describe the same logical input dimension.
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 1},
+      std::vector<xla::DExpr>{input_expr, xla::DExpr::Const(1)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "range",
+                                     std::move(graph), args, &result));
+
+  // The bound is ceil(8 / 3) = 3, while the logical size remains ceil(A / 3).
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_EQ(result_shape.dimensions(0), 3);
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0),
+                                  expected_output_expr));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       ShapeValueChainPreservesExpressionsThroughFill) {
+  // Shape converts [A, 24] into the value vector {A, 24}. Fill then consumes
+  // those values as dimensions, so its output should recover [A, 24].
+  Scope inference_scope = Scope::NewRootScope().ExitOnError();
+  auto placeholder = ops::Placeholder(
+      inference_scope.WithOpName("placeholder"), DT_FLOAT,
+      ops::Placeholder::Shape(PartialTensorShape({-1, 24})));
+  auto fill_shape =
+      ops::Shape(inference_scope.WithOpName("fill_shape"), placeholder);
+  auto fill_value =
+      ops::Const<float>(inference_scope.WithOpName("fill_value"), 1.0f, {});
+  auto inferred_fill = ops::Fill(inference_scope.WithOpName("fill"),
+                                 fill_shape, fill_value);
+
+  grappler::GrapplerItem item;
+  TF_ASSERT_OK(inference_scope.ToGraphDef(&item.graph));
+  grappler::GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(
+      /*assume_valid_feeds=*/false,
+      /*aggressive_shape_inference=*/false,
+      /*include_input_tensor_values=*/false,
+      /*include_output_tensor_values=*/false,
+      /*enable_dynamic_value_inference=*/true));
+
+  const TensorShapeProto& inferred_input_shape =
+      properties.GetOutputProperties("placeholder").at(0).shape();
+  const TensorShapeProto& inferred_fill_shape =
+      properties.GetOutputProperties("fill").at(0).shape();
+  ASSERT_EQ(inferred_input_shape.expressions_size(), 2);
+  ASSERT_EQ(inferred_fill_shape.expressions_size(), 2);
+  const xla::DExpr input_expr =
+      DimExprFromProto(inferred_input_shape.expressions(0));
+  const xla::DExpr expected_output_expr =
+      DimExprFromProto(inferred_fill_shape.expressions(0));
+  EXPECT_TRUE(xla::DynExpr::equal(expected_output_expr, input_expr));
+
+  // Compile the equivalent graph and check both dimensions reach HLO.
+  Scope compilation_scope = Scope::NewRootScope().ExitOnError();
+  auto arg = ops::_Arg(compilation_scope.WithOpName("arg"), DT_FLOAT, 0);
+  auto compiled_fill_shape =
+      ops::Shape(compilation_scope.WithOpName("fill_shape"), arg);
+  auto compiled_fill_value =
+      ops::Const<float>(compilation_scope.WithOpName("fill_value"), 1.0f, {});
+  auto fill = ops::Fill(compilation_scope.WithOpName("fill"),
+                        compiled_fill_shape, compiled_fill_value);
+  auto retval =
+      ops::_Retval(compilation_scope.WithOpName("retval"), fill, 0);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(compilation_scope.ToGraph(graph.get()));
+
+  // Use the same inferred variable at the tf2xla boundary to compare the two
+  // inference paths structurally rather than by variable spelling.
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 24},
+      std::vector<xla::DExpr>{input_expr, xla::DExpr::Const(24)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "fill",
+                                     std::move(graph), args, &result));
+
+  // Compilation keeps the physical [8, 24] bounds and the logical [A, 24]
+  // expressions at the cluster output.
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_EQ(result_shape.dimensions(0), 8);
+  EXPECT_EQ(result_shape.dimensions(1), 24);
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0),
+                                  expected_output_expr));
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1),
+                                  xla::DExpr::Const(24)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ReverseSequencePreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto seq_lens = ops::_Arg(scope.WithOpName("seq_lens"), DT_INT32, 1);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("reverse_sequence", "ReverseSequence")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(seq_lens.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tlen", DT_INT32)
+                   .Attr("batch_dim", 0)
+                   .Attr("seq_dim", 1)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* reverse_sequence = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(reverse_sequence));
+  scope.graph()->AddEdge(input.node(), 0, reverse_sequence, 0);
+  scope.graph()->AddEdge(seq_lens.node(), 0, reverse_sequence, 1);
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(reverse_sequence), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {4, 8},
+      std::vector<xla::DExpr>{xla::DExpr::Var(1), xla::DExpr::Var(2)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {4}, std::vector<xla::DExpr>{xla::DExpr::Var(1)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "reverse_sequence", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(1)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Var(2)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(1)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Var(2)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, UniquePreservesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("unique", "Unique")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("out_idx", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* unique = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(unique));
+  scope.graph()->AddEdge(input.node(), 0, unique, 0);
+
+  auto retval0 =
+      ops::_Retval(scope.WithOpName("retval0"), Output(unique, 0), 0);
+  auto retval1 =
+      ops::_Retval(scope.WithOpName("retval1"), Output(unique, 1), 1);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7}, std::vector<xla::DExpr>{xla::DExpr::Var(3)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "unique",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 2);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(3)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[1].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(3)));
+
+  const xla::Shape& values_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  const xla::Shape& indices_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {1});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(values_shape.expressions(0), xla::DExpr::Var(3)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(indices_shape.expressions(0), xla::DExpr::Var(3)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, DynamicPartitionPreservesPartitionExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto data = ops::_Arg(scope.WithOpName("data"), DT_INT32, 0);
+  auto partitions = ops::_Arg(scope.WithOpName("partitions"), DT_INT32, 1);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("dynamic_partition", "DynamicPartition")
+                   .Input(data.node()->name(), 0, DT_INT32)
+                   .Input(partitions.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("num_partitions", 2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* dynamic_partition = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(dynamic_partition));
+  scope.graph()->AddEdge(data.node(), 0, dynamic_partition, 0);
+  scope.graph()->AddEdge(partitions.node(), 0, dynamic_partition, 1);
+
+  auto retval0 = ops::_Retval(scope.WithOpName("retval0"),
+                              Output(dynamic_partition, 0), 0);
+  auto retval1 = ops::_Retval(scope.WithOpName("retval1"),
+                              Output(dynamic_partition, 1), 1);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6}, std::vector<xla::DExpr>{xla::DExpr::Var(4)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6}, std::vector<xla::DExpr>{xla::DExpr::Var(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "dynamic_partition", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 2);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[1].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(4)));
+
+  const xla::Shape& result0_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  const xla::Shape& result1_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {1});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result0_shape.expressions(0), xla::DExpr::Var(4)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result1_shape.expressions(0), xla::DExpr::Var(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       DynamicPartitionBroadcastPreservesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto data = ops::_Arg(scope.WithOpName("data"), DT_INT32, 0);
+  auto partitions = ops::_Arg(scope.WithOpName("partitions"), DT_INT32, 1);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("dynamic_partition", "DynamicPartition")
+                   .Input(data.node()->name(), 0, DT_INT32)
+                   .Input(partitions.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("num_partitions", 2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* dynamic_partition = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(dynamic_partition));
+  scope.graph()->AddEdge(data.node(), 0, dynamic_partition, 0);
+  scope.graph()->AddEdge(partitions.node(), 0, dynamic_partition, 1);
+
+  auto retval0 = ops::_Retval(scope.WithOpName("retval0"),
+                              Output(dynamic_partition, 0), 0);
+  auto retval1 = ops::_Retval(scope.WithOpName("retval1"),
+                              Output(dynamic_partition, 1), 1);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(40), xla::DExpr::Const(3)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8}, std::vector<xla::DExpr>{xla::DExpr::Var(40)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "dynamic_partition_broadcast",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 2);
+  for (int i = 0; i < 2; ++i) {
+    EXPECT_TRUE(xla::DynExpr::equal(
+        result.outputs[i].shape.get_filled_expression(0), xla::DExpr::Var(40)));
+    EXPECT_TRUE(xla::DynExpr::equal(
+        result.outputs[i].shape.get_filled_expression(1), xla::DExpr::Const(3)));
+    const xla::Shape& out_shape =
+        xla::ShapeUtil::GetSubshape(result.xla_output_shape, {i});
+    EXPECT_TRUE(
+        xla::DynExpr::equal(out_shape.expressions(0), xla::DExpr::Var(40)));
+    EXPECT_TRUE(
+        xla::DynExpr::equal(out_shape.expressions(1), xla::DExpr::Const(3)));
+  }
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       DenseBincountMatrixPreservesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto weights = ops::_Arg(scope.WithOpName("weights"), DT_FLOAT, 1);
+  auto size = ops::Const(scope.WithOpName("size"), 5);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("dense_bincount", "DenseBincount")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(size.node()->name(), 0, DT_INT32)
+                   .Input(weights.node()->name(), 0, DT_FLOAT)
+                   .Attr("Tidx", DT_INT32)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("binary_output", false)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* bincount = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, bincount, 0);
+  scope.graph()->AddEdge(size.node(), 0, bincount, 1);
+  scope.graph()->AddEdge(weights.node(), 0, bincount, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(bincount));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(bincount, 0), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {12, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(50), xla::DExpr::Const(4)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_FLOAT;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {12, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(50), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "dense_bincount", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(50)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+
+  const xla::Shape& out_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(xla::DynExpr::equal(out_shape.expressions(0), xla::DExpr::Var(50)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(out_shape.expressions(1), xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       DynamicStitchEmptyPreservesTrailingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto data0 = ops::_Arg(scope.WithOpName("data0"), DT_INT32, 0);
+  auto data1 = ops::_Arg(scope.WithOpName("data1"), DT_INT32, 1);
+  Tensor empty_indices_tensor(DT_INT32, TensorShape({0}));
+  auto indices0 = ops::Const(scope.WithOpName("indices0"), empty_indices_tensor);
+  auto indices1 = ops::Const(scope.WithOpName("indices1"), empty_indices_tensor);
+
+  NodeDef def;
+  std::vector<NodeDefBuilder::NodeOut> indices_inputs = {
+      {indices0.node()->name(), 0, DT_INT32},
+      {indices1.node()->name(), 0, DT_INT32},
+  };
+  std::vector<NodeDefBuilder::NodeOut> data_inputs = {
+      {data0.node()->name(), 0, DT_INT32},
+      {data1.node()->name(), 0, DT_INT32},
+  };
+  TF_ASSERT_OK(NodeDefBuilder("dynamic_stitch_empty", "DynamicStitch")
+                   .Input(indices_inputs)
+                   .Input(data_inputs)
+                   .Attr("N", 2)
+                   .Attr("T", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* dynamic_stitch = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(indices0.node(), 0, dynamic_stitch, 0);
+  scope.graph()->AddEdge(indices1.node(), 0, dynamic_stitch, 1);
+  scope.graph()->AddEdge(data0.node(), 0, dynamic_stitch, 2);
+  scope.graph()->AddEdge(data1.node(), 0, dynamic_stitch, 3);
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"),
+                             Output(dynamic_stitch, 0), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {0, 7},
+      std::vector<xla::DExpr>{xla::DExpr::Const(0), xla::DExpr::Var(61)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {0, 7},
+      std::vector<xla::DExpr>{xla::DExpr::Const(0), xla::DExpr::Var(61)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "dynamic_stitch_empty", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Const(0)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Var(61)));
+
+  const xla::Shape& out_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(out_shape.expressions(0), xla::DExpr::Const(0)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(out_shape.expressions(1), xla::DExpr::Var(61)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       TensorListPushBackStackPreservesElementExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto element = ops::_Arg(scope.WithOpName("element"), DT_INT32, 0);
+  auto element_shape = ops::Const(scope.WithOpName("element_shape"),
+                                  {-1, 3}, {2});
+  auto max_num_elements = ops::Const(scope.WithOpName("max_num_elements"), 4);
+
+  NodeDef empty_def;
+  TF_ASSERT_OK(NodeDefBuilder("empty_list", "EmptyTensorList")
+                   .Input(element_shape.node()->name(), 0, DT_INT32)
+                   .Input(max_num_elements.node()->name(), 0, DT_INT32)
+                   .Attr("element_dtype", DT_INT32)
+                   .Attr("shape_type", DT_INT32)
+                   .Finalize(&empty_def));
+  absl::Status status;
+  Node* empty_list = scope.graph()->AddNode(empty_def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(element_shape.node(), 0, empty_list, 0);
+  scope.graph()->AddEdge(max_num_elements.node(), 0, empty_list, 1);
+  TF_ASSERT_OK(scope.DoShapeInference(empty_list));
+
+  NodeDef push_def;
+  TF_ASSERT_OK(NodeDefBuilder("push_back", "TensorListPushBack")
+                   .Input(empty_list->name(), 0, DT_VARIANT)
+                   .Input(element.node()->name(), 0, DT_INT32)
+                   .Attr("element_dtype", DT_INT32)
+                   .Finalize(&push_def));
+  Node* push_back = scope.graph()->AddNode(push_def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(empty_list, 0, push_back, 0);
+  scope.graph()->AddEdge(element.node(), 0, push_back, 1);
+  TF_ASSERT_OK(scope.DoShapeInference(push_back));
+
+  NodeDef stack_def;
+  TF_ASSERT_OK(NodeDefBuilder("stack", "TensorListStack")
+                   .Input(push_back->name(), 0, DT_VARIANT)
+                   .Input(element_shape.node()->name(), 0, DT_INT32)
+                   .Attr("element_dtype", DT_INT32)
+                   .Attr("num_elements", 4)
+                   .Finalize(&stack_def));
+  Node* stack = scope.graph()->AddNode(stack_def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(push_back, 0, stack, 0);
+  scope.graph()->AddEdge(element_shape.node(), 0, stack, 1);
+  TF_ASSERT_OK(scope.DoShapeInference(stack));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(stack, 0), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {5, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(70), xla::DExpr::Const(3)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "tensor_list_stack", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Var(70)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(3)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ShapeThenReshapePreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto shape_source = ops::_Arg(scope.WithOpName("shape_source"), DT_INT32, 0);
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 1);
+  auto shape = ops::Shape(scope.WithOpName("shape"), shape_source);
+  auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 7},
+      std::vector<xla::DExpr>{xla::DExpr::Var(41), xla::DExpr::Const(7)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 7},
+      std::vector<xla::DExpr>{xla::DExpr::Var(41), xla::DExpr::Const(7)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "shape_then_reshape", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(41)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(7)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(41)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Const(7)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ZerosLikePreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto zeros = ops::ZerosLike(scope.WithOpName("zeros_like"), input);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), zeros, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {9, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(43), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "zeros_like_exprs", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(43)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, OnesLikePreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto ones = ops::OnesLike(scope.WithOpName("ones_like"), input);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), ones, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {9, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(44), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "ones_like_exprs", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(44)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, MatrixDiagPreservesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("matrix_diag", "MatrixDiag")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* matrix_diag = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(matrix_diag));
+  scope.graph()->AddEdge(input.node(), 0, matrix_diag, 0);
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(matrix_diag), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(45), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "matrix_diag_exprs", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(45)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, WhereBuildsDynamicIndexMatrixShape) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_BOOL, 0);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("where", "Where")
+                   .Input(input.node()->name(), 0, DT_BOOL)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* where = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(where));
+  scope.graph()->AddEdge(input.node(), 0, where, 0);
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(where), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_BOOL;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::PRED, {8, 4, 6},
+      std::vector<xla::DExpr>{xla::DExpr::Var(46), xla::DExpr::Const(4),
+                              xla::DExpr::Var(47)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "where",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_EQ(result_shape.dimensions_size(), 2);
+  EXPECT_EQ(result_shape.dimensions(1), 3);
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Const(3)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, DiagDuplicatesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("diag", "Diag")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* diag = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(diag));
+  scope.graph()->AddEdge(input.node(), 0, diag, 0);
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(diag), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {5}, std::vector<xla::DExpr>{xla::DExpr::Var(42)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "diag",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(42)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Var(42)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(42)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Var(42)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, InTopKPreservesBatchExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto predictions = ops::_Arg(scope.WithOpName("predictions"), DT_FLOAT, 0);
+  auto targets = ops::_Arg(scope.WithOpName("targets"), DT_INT32, 1);
+  auto k = ops::Const<int32>(scope.WithOpName("k"), 3, {});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("in_topk", "InTopKV2")
+                   .Input(predictions.node()->name(), 0, DT_FLOAT)
+                   .Input(targets.node()->name(), 0, DT_INT32)
+                   .Input(k.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* in_topk = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(predictions.node(), 0, in_topk, 0);
+  scope.graph()->AddEdge(targets.node(), 0, in_topk, 1);
+  scope.graph()->AddEdge(k.node(), 0, in_topk, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(in_topk));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(in_topk), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {5, 7},
+      std::vector<xla::DExpr>{xla::DExpr::Var(43), xla::DExpr::Const(7)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {5}, std::vector<xla::DExpr>{xla::DExpr::Var(43)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "in_topk",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(43)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(43)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ReshapeCollapsePreservesSymbolicExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto shape = ops::Const<int32>(scope.WithOpName("shape"), {96}, {1});
+  auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {3, 4, 8},
+      std::vector<xla::DExpr>{xla::DExpr::Var(5), xla::DExpr::Const(4),
+                              xla::DExpr::Const(8)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "reshape_collapse", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(5) * xla::DExpr::Const(32)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ReshapeSplitPreservesSymbolicExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto shape = ops::Const<int32>(scope.WithOpName("shape"), {5, 16}, {2});
+  auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {10, 8},
+      std::vector<xla::DExpr>{xla::DExpr::Var(6), xla::DExpr::Const(8)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "reshape_split", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(6) / xla::DExpr::Const(2)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ReshapeSplitAndCollapsePreservesSymbolicExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto shape = ops::Const<int32>(scope.WithOpName("shape"), {4, 64}, {2});
+  auto reshaped = ops::Reshape(scope.WithOpName("reshape"), input, shape);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), reshaped, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 8, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(7), xla::DExpr::Const(8),
+                              xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "reshape_split_collapse",
+                                     std::move(graph), args, &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(7) / xla::DExpr::Const(2)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, GatherV2PreservesUngatheredExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto params = ops::_Arg(scope.WithOpName("params"), DT_INT32, 0);
+  auto indices = ops::Const<int32>(scope.WithOpName("indices"), {0, 2, 4}, {3});
+  auto axis = ops::Const<int32>(scope.WithOpName("axis"), 1, {});
+  auto gathered =
+      ops::GatherV2(scope.WithOpName("gather"), params, indices, axis);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), gathered, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {9, 7},
+      std::vector<xla::DExpr>{xla::DExpr::Var(8), xla::DExpr::Const(7)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "gather_preserve", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(8)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(8)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, TransposePermutesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto perm = ops::Const<int32>(scope.WithOpName("perm"), {1, 0}, {2});
+  auto transposed = ops::Transpose(scope.WithOpName("transpose"), input, perm);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), transposed, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {5, 7},
+      std::vector<xla::DExpr>{xla::DExpr::Var(9), xla::DExpr::Var(10)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "transpose_exprs", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(10)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Var(9)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ExpandDimsInsertsUnitExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto dim = ops::Const<int32>(scope.WithOpName("dim"), 1, {});
+  auto expanded = ops::ExpandDims(scope.WithOpName("expand"), input, dim);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), expanded, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(11), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "expand_dims_exprs", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(11)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(1)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, SqueezeRemovesUnitExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("squeeze", "Squeeze")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("squeeze_dims", {1})
+                   .Finalize(&def));
+  absl::Status status;
+  Node* squeeze = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, squeeze, 0);
+  TF_ASSERT_OK(scope.DoShapeInference(squeeze));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(squeeze), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 1, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(12), xla::DExpr::Const(1),
+                              xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "squeeze_exprs", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(12)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, SplitPreservesAndDividesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto split_dim = ops::Const<int32>(scope.WithOpName("split_dim"), 0, {});
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto split = ops::Split(scope.WithOpName("split"), split_dim, input, 2);
+  auto retval0 = ops::_Retval(scope.WithOpName("retval0"), split.output[0], 0);
+  auto retval1 = ops::_Retval(scope.WithOpName("retval1"), split.output[1], 1);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(13), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "split_exprs", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(13) / xla::DExpr::Const(2)).simplify();
+  ASSERT_EQ(result.outputs.size(), 2);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[1].shape.get_filled_expression(
+                                      0),
+                                  expected));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, TileScalesExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto multiples =
+      ops::Const<int32>(scope.WithOpName("multiples"), {3, 1}, {2});
+  auto tiled = ops::Tile(scope.WithOpName("tile"), input, multiples);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), tiled, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {4, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(14), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "tile_exprs", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(14) * xla::DExpr::Const(3)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, PackInsertsAxisAndPreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input0 = ops::_Arg(scope.WithOpName("input0"), DT_INT32, 0);
+  auto input1 = ops::_Arg(scope.WithOpName("input1"), DT_INT32, 1);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("pack", "Pack")
+                   .Input({NodeDefBuilder::NodeOut(input0.node()->name(), 0,
+                                                   DT_INT32),
+                           NodeDefBuilder::NodeOut(input1.node()->name(), 0,
+                                                   DT_INT32)})
+                   .Attr("T", DT_INT32)
+                   .Attr("N", 2)
+                   .Attr("axis", 1)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* pack = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input0.node(), 0, pack, 0);
+  scope.graph()->AddEdge(input1.node(), 0, pack, 1);
+  TF_ASSERT_OK(scope.DoShapeInference(pack));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(pack), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6}, std::vector<xla::DExpr>{xla::DExpr::Var(15)});
+  args[1] = args[0];
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "pack",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(15)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(2)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, UnpackRemovesAxisAndPreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("unpack", "Unpack")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("num", 3)
+                   .Attr("axis", 2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* unpack = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, unpack, 0);
+  TF_ASSERT_OK(scope.DoShapeInference(unpack));
+
+  auto retval0 = ops::_Retval(scope.WithOpName("retval0"), Output(unpack, 0), 0);
+  auto retval1 = ops::_Retval(scope.WithOpName("retval1"), Output(unpack, 1), 1);
+  auto retval2 = ops::_Retval(scope.WithOpName("retval2"), Output(unpack, 2), 2);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 4, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(16), xla::DExpr::Const(4),
+                              xla::DExpr::Const(3)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "unpack",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 3);
+  for (int i = 0; i < 3; ++i) {
+    EXPECT_TRUE(xla::DynExpr::equal(
+        result.outputs[i].shape.get_filled_expression(0), xla::DExpr::Var(16)));
+    EXPECT_TRUE(xla::DynExpr::equal(result.outputs[i].shape.get_filled_expression(
+                                        1),
+                                    xla::DExpr::Const(4)));
+  }
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ConcatV2AddsLeadingExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1);
+  auto axis = ops::Const<int32>(scope.WithOpName("axis"), 0, {});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("concat", "ConcatV2")
+                   .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0,
+                                                   DT_INT32),
+                           NodeDefBuilder::NodeOut(rhs.node()->name(), 0,
+                                                   DT_INT32)})
+                   .Input(axis.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tidx", DT_INT32)
+                   .Attr("N", 2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* concat = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(lhs.node(), 0, concat, 0);
+  scope.graph()->AddEdge(rhs.node(), 0, concat, 1);
+  scope.graph()->AddEdge(axis.node(), 0, concat, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(concat));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {5, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(17), xla::DExpr::Const(4)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(18), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "concat",
+                                     std::move(graph), args, &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(17) + xla::DExpr::Var(18)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ConcatAddsLeadingExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto axis = ops::Const<int32>(scope.WithOpName("axis"), 0, {});
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("concat", "Concat")
+                   .Input(axis.node()->name(), 0, DT_INT32)
+                   .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0,
+                                                   DT_INT32),
+                           NodeDefBuilder::NodeOut(rhs.node()->name(), 0,
+                                                   DT_INT32)})
+                   .Attr("T", DT_INT32)
+                   .Attr("N", 2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* concat = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(axis.node(), 0, concat, 0);
+  scope.graph()->AddEdge(lhs.node(), 0, concat, 1);
+  scope.graph()->AddEdge(rhs.node(), 0, concat, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(concat));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {5, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(22), xla::DExpr::Const(4)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(23), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "concat_legacy", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(22) + xla::DExpr::Var(23)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ConcatAddsThreeLeadingExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto mid = ops::_Arg(scope.WithOpName("mid"), DT_INT32, 1);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 2);
+  auto axis = ops::Const<int32>(scope.WithOpName("axis"), 0, {});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("concat0", "ConcatV2")
+                   .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0,
+                                                   DT_INT32),
+                           NodeDefBuilder::NodeOut(mid.node()->name(), 0,
+                                                   DT_INT32)})
+                   .Input(axis.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tidx", DT_INT32)
+                   .Attr("N", 2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* concat0 = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(lhs.node(), 0, concat0, 0);
+  scope.graph()->AddEdge(mid.node(), 0, concat0, 1);
+  scope.graph()->AddEdge(axis.node(), 0, concat0, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(concat0));
+
+  NodeDef def1;
+  TF_ASSERT_OK(NodeDefBuilder("concat1", "ConcatV2")
+                   .Input({NodeDefBuilder::NodeOut(concat0->name(), 0,
+                                                   DT_INT32),
+                           NodeDefBuilder::NodeOut(rhs.node()->name(), 0,
+                                                   DT_INT32)})
+                   .Input(axis.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tidx", DT_INT32)
+                   .Attr("N", 2)
+                   .Finalize(&def1));
+  Node* concat1 = scope.graph()->AddNode(def1, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(concat0, 0, concat1, 0);
+  scope.graph()->AddEdge(rhs.node(), 0, concat1, 1);
+  scope.graph()->AddEdge(axis.node(), 0, concat1, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(concat1));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat1), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(3);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {3, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(31), xla::DExpr::Const(4)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {5, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(32), xla::DExpr::Const(4)});
+  args[2].kind = XlaCompiler::Argument::kParameter;
+  args[2].type = DT_INT32;
+  args[2].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(33), xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "concat_three", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected =
+      (xla::DExpr::Var(31) + xla::DExpr::Var(32) + xla::DExpr::Var(33))
+          .simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  expected));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(0), expected));
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ConcatV2PreservesLeadingExpressionOnInnerAxis) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1);
+  auto axis = ops::Const<int32>(scope.WithOpName("axis"), 1, {});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("concat", "ConcatV2")
+                   .Input({NodeDefBuilder::NodeOut(lhs.node()->name(), 0,
+                                                   DT_INT32),
+                           NodeDefBuilder::NodeOut(rhs.node()->name(), 0,
+                                                   DT_INT32)})
+                   .Input(axis.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tidx", DT_INT32)
+                   .Attr("N", 2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* concat = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(lhs.node(), 0, concat, 0);
+  scope.graph()->AddEdge(rhs.node(), 0, concat, 1);
+  scope.graph()->AddEdge(axis.node(), 0, concat, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(concat));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(concat), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 4},
+      std::vector<xla::DExpr>{xla::DExpr::Var(24), xla::DExpr::Const(4)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(24), xla::DExpr::Const(3)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "concat_inner", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(24)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(7)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, AddPreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1);
+  auto sum = ops::Add(scope.WithOpName("add"), lhs, rhs);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), sum, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(44), xla::DExpr::Const(5)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(44), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(44)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       AddSameRankBroadcastPreservesMappedExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1);
+  auto sum = ops::Add(scope.WithOpName("add"), lhs, rhs);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), sum, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 1, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(45), xla::DExpr::Const(1),
+                              xla::DExpr::Const(3)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 4, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(45), xla::DExpr::Const(4),
+                              xla::DExpr::Const(3)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "add_broadcast", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(45)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(3)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, AddDegenerateBroadcastPreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1);
+  auto sum = ops::Add(scope.WithOpName("add"), lhs, rhs);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), sum, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {1, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Const(1), xla::DExpr::Const(5)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(46), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "add_degenerate", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(46)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       MulSameRankBroadcastPreservesMappedExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_INT32, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_INT32, 1);
+  auto product = ops::Mul(scope.WithOpName("mul"), lhs, rhs);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), product, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 1, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(47), xla::DExpr::Const(1),
+                              xla::DExpr::Const(3)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_INT32;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 4, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Var(47), xla::DExpr::Const(4),
+                              xla::DExpr::Const(3)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "mul_broadcast", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(47)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(3)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, ReverseV2PreservesExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto axis = ops::Const<int32>(scope.WithOpName("axis"), {1}, {1});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("reverse", "ReverseV2")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(axis.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tidx", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* reverse = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, reverse, 0);
+  scope.graph()->AddEdge(axis.node(), 0, reverse, 1);
+  TF_ASSERT_OK(scope.DoShapeInference(reverse));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(reverse), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {8, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(19), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "reverse_v2",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(19)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, BatchMatMulPreservesBatchExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_FLOAT, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_FLOAT, 1);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("batch_matmul", "BatchMatMul")
+                   .Input(lhs.node()->name(), 0, DT_FLOAT)
+                   .Input(rhs.node()->name(), 0, DT_FLOAT)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("adj_x", false)
+                   .Attr("adj_y", false)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* batch_matmul = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(lhs.node(), 0, batch_matmul, 0);
+  scope.graph()->AddEdge(rhs.node(), 0, batch_matmul, 1);
+  TF_ASSERT_OK(scope.DoShapeInference(batch_matmul));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(batch_matmul), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 4, 6},
+      std::vector<xla::DExpr>{xla::DExpr::Var(25), xla::DExpr::Const(4),
+                              xla::DExpr::Const(6)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_FLOAT;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 6, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(25), xla::DExpr::Const(6),
+                              xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "batch_matmul", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(25)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, BatchMatMulV2BroadcastsBatchExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto lhs = ops::_Arg(scope.WithOpName("lhs"), DT_FLOAT, 0);
+  auto rhs = ops::_Arg(scope.WithOpName("rhs"), DT_FLOAT, 1);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("batch_matmul_v2", "BatchMatMulV2")
+                   .Input(lhs.node()->name(), 0, DT_FLOAT)
+                   .Input(rhs.node()->name(), 0, DT_FLOAT)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("adj_x", false)
+                   .Attr("adj_y", false)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* batch_matmul = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(lhs.node(), 0, batch_matmul, 0);
+  scope.graph()->AddEdge(rhs.node(), 0, batch_matmul, 1);
+  TF_ASSERT_OK(scope.DoShapeInference(batch_matmul));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(batch_matmul), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(2);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {1, 4, 6},
+      std::vector<xla::DExpr>{xla::DExpr::Const(1), xla::DExpr::Const(4),
+                              xla::DExpr::Const(6)});
+  args[1].kind = XlaCompiler::Argument::kParameter;
+  args[1].type = DT_FLOAT;
+  args[1].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 6, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(26), xla::DExpr::Const(6),
+                              xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "batch_matmul_v2", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(26)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, SlicePreservesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0, 2}, {2});
+  auto size = ops::Const<int32>(scope.WithOpName("size"), {-1, 3}, {2});
+  auto sliced = ops::Slice(scope.WithOpName("slice"), input, begin, size);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), sliced, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 8},
+      std::vector<xla::DExpr>{xla::DExpr::Var(20), xla::DExpr::Const(8)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "slice",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(20)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(3)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, SliceSubtractsFromLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {2, 1}, {2});
+  auto size = ops::Const<int32>(scope.WithOpName("size"), {-1, 3}, {2});
+  auto sliced = ops::Slice(scope.WithOpName("slice"), input, begin, size);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), sliced, 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {9, 8},
+      std::vector<xla::DExpr>{xla::DExpr::Var(27), xla::DExpr::Const(8)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "slice_subtract", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  (xla::DExpr::Var(27) - 2).simplify()));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(3)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       StridedSlicePreservesLeadingExpressionOnInnerAxis) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0, 1}, {2});
+  auto end = ops::Const<int32>(scope.WithOpName("end"), {0, 7}, {2});
+  auto strides = ops::Const<int32>(scope.WithOpName("strides"), {1, 2}, {2});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(begin.node()->name(), 0, DT_INT32)
+                   .Input(end.node()->name(), 0, DT_INT32)
+                   .Input(strides.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Index", DT_INT32)
+                   .Attr("begin_mask", 1)
+                   .Attr("end_mask", 1)
+                   .Attr("ellipsis_mask", 0)
+                   .Attr("new_axis_mask", 0)
+                   .Attr("shrink_axis_mask", 0)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* strided_slice = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, strided_slice, 0);
+  scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1);
+  scope.graph()->AddEdge(end.node(), 0, strided_slice, 2);
+  scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3);
+  TF_ASSERT_OK(scope.DoShapeInference(strided_slice));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 8},
+      std::vector<xla::DExpr>{xla::DExpr::Var(40), xla::DExpr::Const(8)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "strided_slice_inner", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(40)));
+  EXPECT_EQ(result.outputs[0].shape.dim_size(1), 3);
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, StridedSliceScalesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0, 0}, {2});
+  auto end = ops::Const<int32>(scope.WithOpName("end"), {0, 4}, {2});
+  auto strides = ops::Const<int32>(scope.WithOpName("strides"), {2, 1}, {2});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(begin.node()->name(), 0, DT_INT32)
+                   .Input(end.node()->name(), 0, DT_INT32)
+                   .Input(strides.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Index", DT_INT32)
+                   .Attr("begin_mask", 1)
+                   .Attr("end_mask", 1)
+                   .Attr("ellipsis_mask", 0)
+                   .Attr("new_axis_mask", 0)
+                   .Attr("shrink_axis_mask", 0)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* strided_slice = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, strided_slice, 0);
+  scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1);
+  scope.graph()->AddEdge(end.node(), 0, strided_slice, 2);
+  scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3);
+  TF_ASSERT_OK(scope.DoShapeInference(strided_slice));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 4},
+      std::vector<xla::DExpr>{
+          ((xla::DExpr::Const(2) * xla::DExpr::Var(41)) -
+           xla::DExpr::Const(1))
+              .simplify(),
+          xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "strided_slice_leading", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(
+      result.outputs[0].shape.get_filled_expression(0), xla::DExpr::Var(41)));
+  EXPECT_EQ(result.outputs[0].shape.dim_size(0), 4);
+  EXPECT_EQ(result.outputs[0].shape.dim_size(1), 4);
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, StridedSliceNewAxisInsertsUnitExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0, 0}, {2});
+  auto end = ops::Const<int32>(scope.WithOpName("end"), {0, 0}, {2});
+  auto strides = ops::Const<int32>(scope.WithOpName("strides"), {1, 1}, {2});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(begin.node()->name(), 0, DT_INT32)
+                   .Input(end.node()->name(), 0, DT_INT32)
+                   .Input(strides.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Index", DT_INT32)
+                   .Attr("begin_mask", 0x1)
+                   .Attr("end_mask", 0x1)
+                   .Attr("ellipsis_mask", 0)
+                   .Attr("new_axis_mask", 0x2)
+                   .Attr("shrink_axis_mask", 0)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* strided_slice = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, strided_slice, 0);
+  scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1);
+  scope.graph()->AddEdge(end.node(), 0, strided_slice, 2);
+  scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3);
+  TF_ASSERT_OK(scope.DoShapeInference(strided_slice));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(42), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "strided_slice_new_axis",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(42)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(1)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       StridedSliceShrinkAxisPreservesRemainingExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0, 2, 0}, {3});
+  auto end = ops::Const<int32>(scope.WithOpName("end"), {0, 2, 5}, {3});
+  auto strides = ops::Const<int32>(scope.WithOpName("strides"), {1, 1, 1}, {3});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(begin.node()->name(), 0, DT_INT32)
+                   .Input(end.node()->name(), 0, DT_INT32)
+                   .Input(strides.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Index", DT_INT32)
+                   .Attr("begin_mask", 0x3)
+                   .Attr("end_mask", 0x3)
+                   .Attr("ellipsis_mask", 0)
+                   .Attr("new_axis_mask", 0)
+                   .Attr("shrink_axis_mask", 0x2)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* strided_slice = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, strided_slice, 0);
+  scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1);
+  scope.graph()->AddEdge(end.node(), 0, strided_slice, 2);
+  scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3);
+  TF_ASSERT_OK(scope.DoShapeInference(strided_slice));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 5, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(43), xla::DExpr::Const(5),
+                              xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "strided_slice_shrink", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(43)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       StridedSliceNegativeStridePreservesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0, 0}, {2});
+  auto end = ops::Const<int32>(scope.WithOpName("end"), {0, 0}, {2});
+  auto strides =
+      ops::Const<int32>(scope.WithOpName("strides"), {-1, 1}, {2});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(begin.node()->name(), 0, DT_INT32)
+                   .Input(end.node()->name(), 0, DT_INT32)
+                   .Input(strides.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Index", DT_INT32)
+                   .Attr("begin_mask", 0x3)
+                   .Attr("end_mask", 0x3)
+                   .Attr("ellipsis_mask", 0)
+                   .Attr("new_axis_mask", 0)
+                   .Attr("shrink_axis_mask", 0)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* strided_slice = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, strided_slice, 0);
+  scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1);
+  scope.graph()->AddEdge(end.node(), 0, strided_slice, 2);
+  scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3);
+  TF_ASSERT_OK(scope.DoShapeInference(strided_slice));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(44), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "strided_slice_negative",
+                                     std::move(graph), args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(44)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest,
+       StridedSliceNegativeStrideTwoScalesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto begin = ops::Const<int32>(scope.WithOpName("begin"), {0, 0}, {2});
+  auto end = ops::Const<int32>(scope.WithOpName("end"), {0, 0}, {2});
+  auto strides =
+      ops::Const<int32>(scope.WithOpName("strides"), {-2, 1}, {2});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("strided_slice", "StridedSlice")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(begin.node()->name(), 0, DT_INT32)
+                   .Input(end.node()->name(), 0, DT_INT32)
+                   .Input(strides.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Index", DT_INT32)
+                   .Attr("begin_mask", 0x3)
+                   .Attr("end_mask", 0x3)
+                   .Attr("ellipsis_mask", 0)
+                   .Attr("new_axis_mask", 0)
+                   .Attr("shrink_axis_mask", 0)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* strided_slice = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, strided_slice, 0);
+  scope.graph()->AddEdge(begin.node(), 0, strided_slice, 1);
+  scope.graph()->AddEdge(end.node(), 0, strided_slice, 2);
+  scope.graph()->AddEdge(strides.node(), 0, strided_slice, 3);
+  TF_ASSERT_OK(scope.DoShapeInference(strided_slice));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(strided_slice), 0);
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(45), xla::DExpr::Const(5)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "strided_slice_negative_two",
+                                     std::move(graph), args, &result));
 
-class DummyReadResourceOp : public XlaOpKernel {
- public:
-  explicit DummyReadResourceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
-  void Compile(XlaOpKernelContext* ctx) override {
-    ResourceMgr* rm = ctx->op_kernel_context()->resource_manager();
-    OP_REQUIRES(ctx, rm, errors::Internal("No resource manager."));
-    DummyResourceForTest* dummy;
-    OP_REQUIRES_OK(ctx, rm->Lookup<DummyResourceForTest>(
-                            rm->default_container(), "dummy", &dummy));
-    dummy->Increment();
-    dummy->Unref();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(
+      result.outputs[0].shape.get_filled_expression(0),
+      ((xla::DExpr::Var(45) + xla::DExpr::Const(1)) / xla::DExpr::Const(2))
+          .simplify()));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
 
-    ctx->SetOutput(0, ctx->Input(0));
-    ctx->SetOutput(1, ctx->Input(0));
-  }
-};
+TEST_F(XlaCompilerDynamicSizesTest, PadAddsToLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto paddings =
+      ops::Const<int32>(scope.WithOpName("paddings"), {1, 2, 0, 0}, {2, 2});
+  auto padded = ops::Pad(scope.WithOpName("pad"), input, paddings);
+  auto retval = ops::_Retval(scope.WithOpName("retval"), padded, 0);
 
-class DummyReadResourceCC {
- public:
-  DummyReadResourceCC(const Scope& scope, const Input& value) {
-    if (!scope.ok()) return;
-    auto _value = ops::AsNodeOut(scope, value);
-    if (!scope.ok()) return;
-    Node* ret;
-    const auto unique_name = scope.GetUniqueNameForOp("DummyReadResource");
-    auto builder = NodeBuilder(unique_name, "DummyReadResource").Input(_value);
-    scope.UpdateBuilder(&builder);
-    scope.UpdateStatus(builder.Finalize(scope.graph(), &ret));
-    if (!scope.ok()) return;
-    scope.UpdateStatus(scope.DoShapeInference(ret));
-    if (!scope.ok()) return;
-    this->output1_ = Output(ret, 0);
-    this->output2_ = Output(ret, 1);
-  }
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
 
-  Output output1_;
-  Output output2_;
-};
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {7, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(28), xla::DExpr::Const(5)});
 
-REGISTER_OP("DummyReadResource")
-    .Input("input: int32")
-    .Output("output1: int32")
-    .Output("output2: int32")
-    .SetShapeFn(shape_inference::UnknownShape)
-    .Doc(R"doc(
-A dummy Op.
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "pad",
+                                     std::move(graph), args, &result));
 
-input: dummy input.
-output1: dummy output.
-output2: dummy output.
-)doc");
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  (xla::DExpr::Var(28) + 3).simplify()));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
+}
 
-REGISTER_XLA_OP(Name("DummyReadResource"), DummyReadResourceOp);
+TEST_F(XlaCompilerDynamicSizesTest, SpaceToBatchNDScalesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0);
+  auto block_shape =
+      ops::Const<int32>(scope.WithOpName("block_shape"), {2}, {1});
+  auto paddings =
+      ops::Const<int32>(scope.WithOpName("paddings"), {0, 0}, {1, 2});
 
-// DummyDuplicateOp is present purely to test multiple REGISTER_XLA_OP calls
-// on the same Op name below.
-class DummyDuplicateOp : public XlaOpKernel {
- public:
-  explicit DummyDuplicateOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
-  void Compile(XlaOpKernelContext* ctx) override {
-    ctx->SetOutput(0, ctx->Input(0));
-  }
-};
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("space_to_batch", "SpaceToBatchND")
+                   .Input(input.node()->name(), 0, DT_FLOAT)
+                   .Input(block_shape.node()->name(), 0, DT_INT32)
+                   .Input(paddings.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("Tblock_shape", DT_INT32)
+                   .Attr("Tpaddings", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* space_to_batch = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, space_to_batch, 0);
+  scope.graph()->AddEdge(block_shape.node(), 0, space_to_batch, 1);
+  scope.graph()->AddEdge(paddings.node(), 0, space_to_batch, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(space_to_batch));
 
-REGISTER_OP("DummyDuplicateOp")
-    .Input("input: int32")
-    .Output("output: int32")
-    .Doc(R"doc(
-A dummy Op.
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(space_to_batch), 0);
 
-input: dummy input.
-output: dummy output.
-)doc");
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
 
-REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_CPU_XLA_JIT),
-                DummyDuplicateOp);
-REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_GPU_XLA_JIT),
-                DummyDuplicateOp);
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {4, 8},
+      std::vector<xla::DExpr>{xla::DExpr::Var(29), xla::DExpr::Const(8)});
 
-// Tests compilation and execution of an empty graph.
-TEST_F(XlaCompilerTest, EmptyReturnValues) {
   XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "space_to_batch", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  (xla::DExpr::Const(2) * xla::DExpr::Var(29))
+                                      .simplify()));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+}
+
+TEST_F(XlaCompilerDynamicSizesTest, BatchToSpaceNDDividesLeadingExpression) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0);
+  auto block_shape =
+      ops::Const<int32>(scope.WithOpName("block_shape"), {2}, {1});
+  auto crops = ops::Const<int32>(scope.WithOpName("crops"), {0, 0}, {1, 2});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("batch_to_space", "BatchToSpaceND")
+                   .Input(input.node()->name(), 0, DT_FLOAT)
+                   .Input(block_shape.node()->name(), 0, DT_INT32)
+                   .Input(crops.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("Tblock_shape", DT_INT32)
+                   .Attr("Tcrops", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* batch_to_space = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, batch_to_space, 0);
+  scope.graph()->AddEdge(block_shape.node(), 0, batch_to_space, 1);
+  scope.graph()->AddEdge(crops.node(), 0, batch_to_space, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(batch_to_space));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(batch_to_space), 0);
 
   std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
-  XlaCompiler::CompilationResult result;
-  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add",
-                                     std::move(graph),
-                                     /*args=*/{}, &result));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
 
-  TF_ASSERT_OK(client_->Execute(*result.computation, {}).status());
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {8, 4},
+      std::vector<xla::DExpr>{(xla::DExpr::Const(2) * xla::DExpr::Var(30))
+                                  .simplify(),
+                              xla::DExpr::Const(4)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "batch_to_space", std::move(graph), args,
+                                     &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(30)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(8)));
 }
 
-// Tests compilation and execution of a graph that adds two tensors.
-TEST_F(XlaCompilerTest, Simple) {
-  // Builds a graph that adds two Tensors.
+TEST_F(XlaCompilerDynamicSizesTest, SpaceToDepthScalesDepthExpression) {
   Scope scope = Scope::NewRootScope().ExitOnError();
-  auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
-  auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
-  auto c = ops::Add(scope.WithOpName("C"), a, b);
-  auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0);
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("space_to_depth", "SpaceToDepth")
+                   .Input(input.node()->name(), 0, DT_FLOAT)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("block_size", 2)
+                   .Attr("data_format", "NHWC")
+                   .Finalize(&def));
+  absl::Status status;
+  Node* space_to_depth = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, space_to_depth, 0);
+  TF_ASSERT_OK(scope.DoShapeInference(space_to_depth));
+
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(space_to_depth), 0);
+
   std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
   TF_ASSERT_OK(scope.ToGraph(graph.get()));
 
-  // Builds a description of the arguments.
-  std::vector<XlaCompiler::Argument> args(2);
+  std::vector<XlaCompiler::Argument> args(1);
   args[0].kind = XlaCompiler::Argument::kParameter;
-  args[0].type = DT_INT32;
-  args[0].shape = TensorShape({2});
-  args[1].kind = XlaCompiler::Argument::kParameter;
-  args[1].type = DT_INT32;
-  args[1].shape = TensorShape({2});
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {5, 8, 8, 3},
+      std::vector<xla::DExpr>{xla::DExpr::Const(5), xla::DExpr::Const(8),
+                              xla::DExpr::Const(8), xla::DExpr::Var(35)});
 
-  // Compiles the graph.
   XlaCompiler compiler(DefaultOptions());
 
   XlaCompiler::CompilationResult result;
-  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add",
-                                     std::move(graph), args, &result));
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "space_to_depth", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected_depth =
+      (xla::DExpr::Const(4) * xla::DExpr::Var(35)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Const(5)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      3),
+                                  expected_depth));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Const(5)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(1), xla::DExpr::Const(4)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(2), xla::DExpr::Const(4)));
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(3), expected_depth));
+}
 
-  // Tests that the generated computation works.
-  xla::Literal param0_literal = xla::LiteralUtil::CreateR1<int32>({7, 42});
-  xla::Literal param1_literal = xla::LiteralUtil::CreateR1<int32>({-3, 101});
-  std::unique_ptr<xla::GlobalData> param0_data =
-      client_->TransferToServer(param0_literal).value();
-  std::unique_ptr<xla::GlobalData> param1_data =
-      client_->TransferToServer(param1_literal).value();
+TEST_F(XlaCompilerDynamicSizesTest, DepthToSpaceScalesSpatialExpressions) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_FLOAT, 0);
 
-  std::unique_ptr<xla::GlobalData> actual =
-      client_
-          ->Execute(*result.computation, {param0_data.get(), param1_data.get()})
-          .value();
-  xla::Literal actual_literal = client_->Transfer(*actual).value();
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("depth_to_space", "DepthToSpace")
+                   .Input(input.node()->name(), 0, DT_FLOAT)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("block_size", 2)
+                   .Attr("data_format", "NHWC")
+                   .Finalize(&def));
+  absl::Status status;
+  Node* depth_to_space = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, depth_to_space, 0);
+  TF_ASSERT_OK(scope.DoShapeInference(depth_to_space));
 
-  xla::Literal expected0 = xla::LiteralUtil::CreateR1<int32>({4, 143});
-  xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0});
-  EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal));
-}
+  auto retval =
+      ops::_Retval(scope.WithOpName("retval"), Output(depth_to_space), 0);
 
-absl::StatusOr<std::unique_ptr<xla::HloModule>> LoadModuleFromHloProto(
-    const xla::HloModuleProto& module_proto) {
-  TF_ASSIGN_OR_RETURN(auto module_config,
-                      xla::HloModule::CreateModuleConfigFromProto(
-                          module_proto, xla::GetDebugOptionsFromFlags()));
-  return xla::CreateModuleFromProto(module_proto, module_config);
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::F32, {5, 4, 4, 12},
+      std::vector<xla::DExpr>{xla::DExpr::Const(5), xla::DExpr::Var(37),
+                              xla::DExpr::Const(4), xla::DExpr::Const(12)});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "depth_to_space", std::move(graph), args,
+                                     &result));
+
+  xla::DExpr expected_height =
+      (xla::DExpr::Const(2) * xla::DExpr::Var(37)).simplify();
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Const(5)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  expected_height));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      2),
+                                  xla::DExpr::Const(8)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      3),
+                                  xla::DExpr::Const(3)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Const(5)));
+  EXPECT_TRUE(xla::DynExpr::equal(result_shape.expressions(1), expected_height));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(2), xla::DExpr::Const(8)));
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(3), xla::DExpr::Const(3)));
 }
 
-// Tests compilation and execution of a graph that adds two tensors with dynamic
-// shape parameters.
-TEST_F(XlaCompilerTest, SimpleDynamicShapeParameter) {
-  // Builds a graph that adds two Tensors.
+TEST_F(XlaCompilerDynamicSizesTest, RollPreservesExpressions) {
   Scope scope = Scope::NewRootScope().ExitOnError();
-  auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
-  auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
-  auto c = ops::Add(scope.WithOpName("C"), a, b);
-  auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
+  auto input = ops::_Arg(scope.WithOpName("input"), DT_INT32, 0);
+  auto shift = ops::Const<int32>(scope.WithOpName("shift"), 2, {});
+  auto axis = ops::Const<int32>(scope.WithOpName("axis"), 1, {});
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("roll", "Roll")
+                   .Input(input.node()->name(), 0, DT_INT32)
+                   .Input(shift.node()->name(), 0, DT_INT32)
+                   .Input(axis.node()->name(), 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tshift", DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* roll = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  scope.graph()->AddEdge(input.node(), 0, roll, 0);
+  scope.graph()->AddEdge(shift.node(), 0, roll, 1);
+  scope.graph()->AddEdge(axis.node(), 0, roll, 2);
+  TF_ASSERT_OK(scope.DoShapeInference(roll));
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(roll), 0);
+
   std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
   TF_ASSERT_OK(scope.ToGraph(graph.get()));
 
-  // Builds a description of the arguments.
-  std::vector<XlaCompiler::Argument> args(2);
+  std::vector<XlaCompiler::Argument> args(1);
   args[0].kind = XlaCompiler::Argument::kParameter;
   args[0].type = DT_INT32;
-  args[0].shape =
-      xla::ShapeUtil::MakeShape(/*element_type=*/xla::S32, /*dimensions=*/{2},
-                                /*dynamic_dimensions=*/{true});
-  args[1].kind = XlaCompiler::Argument::kParameter;
-  args[1].type = DT_INT32;
-  args[1].shape = TensorShape(/*dimensions=*/{2});
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {6, 5},
+      std::vector<xla::DExpr>{xla::DExpr::Var(21), xla::DExpr::Const(5)});
 
-  // Compiles the graph.
   XlaCompiler compiler(DefaultOptions());
 
   XlaCompiler::CompilationResult result;
-  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add",
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "roll",
                                      std::move(graph), args, &result));
 
-  auto hlo = result.computation->proto();
-  TF_ASSERT_OK_AND_ASSIGN(auto module, LoadModuleFromHloProto(hlo));
-  EXPECT_EQ(module->computation_count(), 1);
-  EXPECT_TRUE(module->mutable_computation(0)
-                  ->parameter_instruction(0)
-                  ->shape()
-                  .is_dynamic());
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(21)));
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      1),
+                                  xla::DExpr::Const(5)));
 }
 
 // Tests compilation of a graph where the _Retval node is not necessarily last
@@ -1014,6 +3862,13 @@ FunctionDef FillFn() {
       {{{"y"}, "Fill", {"dims", "x"}, {{"T", "$T"}}}});
 }
 
+FunctionDef IdentityFn() {
+  return FunctionDefHelper::Define(
+      "IdentityFn", {"x: T"}, {"y: T"},
+      {"T: {float, double, int32, int64}"},
+      {{{"y"}, "Identity", {"x"}, {{"T", "$T"}}}});
+}
+
 TEST_F(XlaCompilerTest, FunctionCallWithConstants) {
   // Certain operations in a function, "Fill" for example, requires the
   // operator's argument to be a compile-time constant instead of a parameter.
@@ -1057,6 +3912,53 @@ TEST_F(XlaCompilerTest, FunctionCallWithConstants) {
                                      std::move(graph), args, &result));
 }
 
+TEST_F(XlaCompilerDynamicSizesTest, FunctionCallPreservesDynamicExpressions) {
+  XlaCompiler compiler(DefaultOptions());
+
+  FunctionDefLibrary flib;
+  *flib.add_function() = IdentityFn();
+  TF_ASSERT_OK(flib_def_->AddFunctionDef(IdentityFn()));
+
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto arg = ops::_Arg(scope.WithOpName("arg"), DT_INT32, 0);
+  TF_EXPECT_OK(scope.graph()->AddFunctionLibrary(flib));
+
+  NodeDef def;
+  TF_ASSERT_OK(NodeDefBuilder("identity_fn", "IdentityFn", flib_def_.get())
+                   .Input(arg.node()->name(), 0, DT_INT32)
+                   .Finalize(&def));
+  absl::Status status;
+  Node* identity_fn = scope.graph()->AddNode(def, &status);
+  TF_ASSERT_OK(status);
+  TF_ASSERT_OK(scope.DoShapeInference(identity_fn));
+  scope.graph()->AddEdge(arg.node(), 0, identity_fn, 0);
+
+  auto retval = ops::_Retval(scope.WithOpName("retval"), Output(identity_fn), 0);
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_INT32;
+  args[0].shape = xla::ShapeUtil::MakeShape(
+      xla::S32, {9}, std::vector<xla::DExpr>{xla::DExpr::Var(2)});
+
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(),
+                                     "identity_function", std::move(graph),
+                                     args, &result));
+
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_TRUE(xla::DynExpr::equal(result.outputs[0].shape.get_filled_expression(
+                                      0),
+                                  xla::DExpr::Var(2)));
+
+  const xla::Shape& result_shape =
+      xla::ShapeUtil::GetSubshape(result.xla_output_shape, {0});
+  EXPECT_TRUE(
+      xla::DynExpr::equal(result_shape.expressions(0), xla::DExpr::Var(2)));
+}
+
 // Tests CompileFunction with a local function lookup failing, fails with
 // informative error about both lookups.
 TEST_F(XlaCompilerTest, LocalFunctionWithWrongArgumentsFail) {
diff --git a/tensorflow/compiler/tf2xla/xla_expression.cc b/tensorflow/compiler/tf2xla/xla_expression.cc
index 61bd10e413ccf3..0488a5170ca211 100644
--- a/tensorflow/compiler/tf2xla/xla_expression.cc
+++ b/tensorflow/compiler/tf2xla/xla_expression.cc
@@ -17,9 +17,12 @@ limitations under the License.
 
 #include "tensorflow/compiler/tf2xla/literal_util.h"
 #include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/compiler/tf2xla/symbolic_content_util.h"
 #include "xla/hlo/builder/value_inference.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/logging.h"
 
 namespace tensorflow {
 
@@ -73,6 +76,50 @@ XlaExpression XlaExpression::Resource(XlaResource* resource) {
   return e;
 }
 
+void XlaExpression::set_contents(std::vector<xla::DExpr> contents) {
+  switch (kind_) {
+    case Kind::kXlaOp:
+    case Kind::kTensorList:
+      if (handle_.valid() && !handle_.IsUninitialized()) {
+        if (SymbolicContentEnabled()) {
+          auto status = handle_.builder()->SetInstructionContents(
+              handle_, std::move(contents));
+          if (!status.ok()) {
+            VLOG(1) << "Failed to set XlaOp contents: " << status;
+          }
+          return;
+        }
+        break;
+      }
+      break;
+    case Kind::kInvalid:
+    case Kind::kConstant:
+    case Kind::kResource:
+      break;
+  }
+  local_contents_ = std::move(contents);
+}
+
+absl::Span<const xla::DExpr> XlaExpression::contents() const {
+  switch (kind_) {
+    case Kind::kXlaOp:
+    case Kind::kTensorList:
+      if (SymbolicContentEnabled() && handle_.valid() &&
+          !handle_.IsUninitialized()) {
+        auto contents_or = handle_.builder()->GetInstructionContents(handle_);
+        if (contents_or.ok()) {
+          return **contents_or;
+        }
+      }
+      break;
+    case Kind::kInvalid:
+    case Kind::kConstant:
+    case Kind::kResource:
+      break;
+  }
+  return local_contents_;
+}
+
 string XlaExpression::HumanString() const {
   switch (kind_) {
     case Kind::kInvalid:
@@ -95,7 +142,11 @@ xla::XlaOp XlaExpression::AsXlaOp(xla::XlaBuilder* builder) const {
         xla::BorrowingLiteral literal;
         TF_RETURN_IF_ERROR(
             HostTensorToBorrowingLiteral(*constant_value_, &literal));
-        return xla::ConstantLiteral(builder, literal);
+        xla::XlaOp op = xla::ConstantLiteral(builder, literal);
+        if (!local_contents_.empty()) {
+          TF_RETURN_IF_ERROR(builder->SetInstructionContents(op, local_contents_));
+        }
+        return op;
       }
       case Kind::kTensorList:
         TF_FALLTHROUGH_INTENDED;
diff --git a/tensorflow/compiler/tf2xla/xla_expression.h b/tensorflow/compiler/tf2xla/xla_expression.h
index d410b79a3da137..980ab6dba4b73a 100644
--- a/tensorflow/compiler/tf2xla/xla_expression.h
+++ b/tensorflow/compiler/tf2xla/xla_expression.h
@@ -16,11 +16,15 @@ limitations under the License.
 #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_
 #define TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_
 
+#include <vector>
+
 #include "absl/types/optional.h"
+#include "absl/types/span.h"
 #include "tensorflow/compiler/tf2xla/xla_resource.h"
 #include "xla/client/client.h"
 #include "xla/hlo/builder/value_inference.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/statusor.h"
@@ -112,6 +116,13 @@ class XlaExpression {
   // Return the dynamism of the expression, if available.
   std::optional<Tensor> value_dynamism() const { return value_dynamism_; }
 
+  // Set symbolic content metadata for expressions whose values should retain
+  // links to symbolic dimensions across shape-tensor flows.
+  void set_contents(std::vector<xla::DExpr> contents);
+
+  // Return symbolic content metadata.
+  absl::Span<const xla::DExpr> contents() const;
+
   XlaResource* resource() const { return resource_; }
 
   // Returns a human-readable summary of the expression.
@@ -164,6 +175,12 @@ class XlaExpression {
   // Indicate whether each value inside a tensor is dynamic or not.
   std::optional<Tensor> value_dynamism_;
 
+  // For constant expressions, marks a single element that later passes may
+  // reinterpret as coming from a dynamic expression instead of the literal.
+  // Symbolic expressions describing tensor contents when this expression is
+  // used as a shape-like value.
+  std::vector<xla::DExpr> local_contents_;
+
   // The resource, if kind_ == kResource. Not owned.
   XlaResource* resource_ = nullptr;
 };
diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc
index 4a570827029330..a1dce933db5309 100644
--- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc
+++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc
@@ -464,6 +464,12 @@ absl::Status XlaOpKernelContext::ConstantInputAsShape(
           ", result: ", num_elements);
   }
   *shape = TensorShape(dims);
+  const auto& contents = InputExpression(index).contents();
+  for (int i = 0; i < shape->dims() && i < contents.size(); ++i) {
+    if (contents[i] && contents[i]->is_dynamic()) {
+      shape->set_expression(i, contents[i]);
+    }
+  }
   return absl::OkStatus();
 }
 
@@ -736,7 +742,8 @@ absl::Status AssignVariableTensor(const Tensor& tensor, DataType type,
   xla::Shape xla_shape;
   TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape));
   if (!xla::ShapeUtil::Compatible(xla_shape, representation_shape)) {
-    handle = xla::Reshape(handle, representation_shape.dimensions());
+    handle = xla::Reshape(handle, representation_shape.dimensions(),
+                          representation_shape.expressions());
   }
   variable->SetRepresentationShape(representation_shape);
   return variable->SetValue(handle);
diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc
index 445065971f2a6a..92493928bc2d6f 100644
--- a/tensorflow/compiler/tf2xla/xla_op_registry.cc
+++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc
@@ -415,6 +415,18 @@ std::vector<const KernelDef*> XlaOpRegistry::DeviceKernels(
   return ops;
 }
 
+/* static */ bool XlaOpRegistry::IsMlirXlaOp(absl::string_view op) {
+  XlaOpRegistry& registry = Instance();
+  mutex_lock lock(registry.mutex_);
+  auto it = registry.ops_.find(string(op));
+  if (it == registry.ops_.end() || it->second.empty()) {
+    return false;
+  }
+  return absl::c_all_of(it->second, [](const auto& registration) {
+    return registration->uses_mlir_kernel;
+  });
+}
+
 /*static*/ const std::unordered_set<std::string>*
 XlaOpRegistry::CompileTimeConstantInputArgNames(const string& op) {
   XlaOpRegistry& registry = Instance();
@@ -604,8 +616,9 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::Label(std::string label) {
 }
 
 std::unique_ptr<XlaOpRegistry::OpRegistration> XlaOpRegistrationBuilder::Build(
-    XlaOpRegistry::Factory factory) {
+    XlaOpRegistry::Factory factory, bool uses_mlir_kernel) {
   registration_->factory = factory;
+  registration_->uses_mlir_kernel = uses_mlir_kernel;
   return std::move(registration_);
 }
 
diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.h b/tensorflow/compiler/tf2xla/xla_op_registry.h
index 5eaf0fb2d42bfa..b6abc2a39f71f2 100644
--- a/tensorflow/compiler/tf2xla/xla_op_registry.h
+++ b/tensorflow/compiler/tf2xla/xla_op_registry.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include <functional>
 #include <memory>
 #include <set>
+#include <type_traits>
 #include <unordered_map>
 #include <vector>
 
@@ -44,6 +45,8 @@ limitations under the License.
 
 namespace tensorflow {
 
+class MlirXlaOpKernel;
+
 // Names of the XLA compilation devices. These are not user-visible, and are
 // used internally by the Tensorflow/XLA bridge to perform symbolic execution of
 // a Tensorflow graph.
@@ -233,6 +236,10 @@ class XlaOpRegistry {
   // Returns all operations for which there are XLA kernels on any device.
   static std::vector<string> GetAllRegisteredOps();
 
+  // Returns true if the operation is lowered exclusively through
+  // MlirXlaOpKernel.
+  static bool IsMlirXlaOp(absl::string_view op);
+
   // Returns (via `result`) the indices of inputs to `node_def` that must be
   // compile-time constants. Returns an empty vector if the op is not
   // registered.
@@ -339,6 +346,8 @@ class XlaOpRegistry {
     // operands and not their values.
     bool is_metadata_op = false;
 
+    bool uses_mlir_kernel = false;
+
     std::string label;
 
     // Factory used to build OpKernels that perform symbolic execution.
@@ -383,6 +392,9 @@ class XlaOpRegistry {
 #define REGISTER_XLA_OP(NAME, OP) \
   REGISTER_XLA_OP_UNIQ_HELPER(__COUNTER__, NAME, OP)
 
+#define REGISTER_XLA_OP_FACTORY(NAME, FACTORY) \
+  REGISTER_XLA_OP_FACTORY_UNIQ_HELPER(__COUNTER__, NAME, FACTORY)
+
 #define REGISTER_XLA_CONV_OP(BUILDER, OP)                                      \
   REGISTER_XLA_OP(BUILDER.TypeConstraint("T", GetXlaConvTypesForNonGpu()), OP) \
   REGISTER_XLA_OP(BUILDER.TypeConstraint("T", GetXlaConvTypesForGpu())         \
@@ -431,7 +443,7 @@ class XlaOpRegistrationBuilder {
   XlaOpRegistrationBuilder& Label(std::string label);
 
   std::unique_ptr<XlaOpRegistry::OpRegistration> Build(
-      XlaOpRegistry::Factory factory);
+      XlaOpRegistry::Factory factory, bool uses_mlir_kernel = false);
 
  private:
   XlaOpRegistrationBuilder(absl::string_view name);
@@ -454,11 +466,19 @@ class XlaOpRegistrar {
 #define REGISTER_XLA_OP_UNIQ_HELPER(COUNTER, BUILDER, OP) \
   REGISTER_XLA_OP_UNIQ(COUNTER, BUILDER, OP)
 
+#define REGISTER_XLA_OP_FACTORY_UNIQ_HELPER(COUNTER, BUILDER, FACTORY) \
+  REGISTER_XLA_OP_FACTORY_UNIQ(COUNTER, BUILDER, FACTORY)
+
 #define REGISTER_XLA_OP_UNIQ(CTR, BUILDER, OP)                                 \
   static ::tensorflow::XlaOpRegistrar xla_op_registrar__body__##CTR##__object( \
       ::tensorflow::XlaOpRegistrationBuilder::BUILDER.Build(                   \
           [](::tensorflow::OpKernelConstruction* context)                      \
-              -> ::tensorflow::OpKernel* { return new OP(context); }));
+              -> ::tensorflow::OpKernel* { return new OP(context); },           \
+          std::is_same_v<OP, ::tensorflow::MlirXlaOpKernel>));
+
+#define REGISTER_XLA_OP_FACTORY_UNIQ(CTR, BUILDER, FACTORY)                    \
+  static ::tensorflow::XlaOpRegistrar xla_op_registrar__body__##CTR##__object( \
+      ::tensorflow::XlaOpRegistrationBuilder::BUILDER.Build(FACTORY));
 
 class XlaBackendRegistrar {
  public:
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 2f0ff5e91867f1..72bd6f712672e3 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -1740,6 +1740,7 @@ tf_cuda_library(
         "@local_xla//xla/tsl/framework:cancellation",
         "@local_xla//xla/tsl/util:command_line_flags",
         "@local_xla//xla/tsl/util:device_name_utils",
+        "@local_xla//xla:shape_util",
     ] + if_cuda([
         "@local_config_cuda//cuda:cudnn_header",
     ]) + if_static(
diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD
index 301015eba61fe8..f99259ad3e3862 100644
--- a/tensorflow/core/common_runtime/BUILD
+++ b/tensorflow/core/common_runtime/BUILD
@@ -2754,6 +2754,7 @@ tf_cc_test(
         ":direct_session_internal",
         "//tensorflow/cc:cc_ops",
         "//tensorflow/cc:cc_ops_internal",
+        "//tensorflow/cc:function_ops",
         "//tensorflow/cc:sendrecv_ops",
         "//tensorflow/core:framework",
         "//tensorflow/core:framework_internal",
@@ -2769,9 +2770,13 @@ tf_cc_test(
         "//tensorflow/core/kernels:cast_op",
         "//tensorflow/core/kernels:concat_op",
         "//tensorflow/core/kernels:cwise_op",
+        "//tensorflow/core/kernels:gather_op",
         "//tensorflow/core/kernels:identity_op",
         "//tensorflow/core/kernels:immutable_constant_op",
         "//tensorflow/core/kernels:matmul_op",
+        "//tensorflow/core/kernels:pack_op",
+        "//tensorflow/core/kernels:reshape_op",
+        "//tensorflow/core/kernels:slice_op",
         "//tensorflow/core/kernels:topk_op",
         "@eigen_archive//:eigen3",
     ],
diff --git a/tensorflow/core/common_runtime/constant_folding.cc b/tensorflow/core/common_runtime/constant_folding.cc
index 6820a5ddd696d3..9a6f7dc7269ed3 100644
--- a/tensorflow/core/common_runtime/constant_folding.cc
+++ b/tensorflow/core/common_runtime/constant_folding.cc
@@ -31,6 +31,8 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/rendezvous_mgr.h"
 #include "tensorflow/core/framework/log_memory.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/graph/algorithm.h"
@@ -40,6 +42,7 @@ limitations under the License.
 #include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/gtl/flatset.h"
 #include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/denormal.h"
 #include "tensorflow/core/platform/setround.h"
 #include "tensorflow/core/public/session_options.h"
@@ -49,6 +52,461 @@ namespace tensorflow {
 namespace {
 
 const char kScopedAllocatorAttrName[] = "_scoped_allocator";
+const char kXlaShapeDerivedAttrName[] = "_xla_shape_derived";
+const char kUserInferredValueContentsAttrName[] =
+    "_user_inferred_value_contents";
+
+bool IsShapeOp(const Node* n);
+
+bool GetShapeFromArgNode(const Node* node, TensorShapeProto* out_shape) {
+  for (const Edge* edge : node->in_edges()) {
+    if (edge->IsControlEdge()) continue;
+
+    Node* input_node = edge->src();
+    if (input_node->type_string() == "_Arg") {
+      std::vector<TensorShapeProto> shapes;
+      if (GetNodeAttr(input_node->def(), "_output_shapes", &shapes).ok() &&
+          !shapes.empty() && HasDynamicDimExprs(shapes[0])) {
+        *out_shape = shapes[0];
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+bool GetShapeFromDirectDynamicSource(const Node* node,
+                                     TensorShapeProto* out_shape) {
+  return (IsShapeOp(node) ||
+          node->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) &&
+         GetShapeFromArgNode(node, out_shape);
+}
+
+bool TryGetContentsProtoAttr(const AttrSlice& attrs,
+                             TensorShapeProto* out_contents) {
+  string serialized_contents;
+  if (!GetNodeAttr(attrs, kUserInferredValueContentsAttrName,
+                   &serialized_contents)
+           .ok()) {
+    return false;
+  }
+  out_contents->Clear();
+  return out_contents->ParseFromString(serialized_contents);
+}
+
+bool HasTransitiveDynamicShapeContents(
+    const Node* node, std::unordered_map<const Node*, bool>* memo,
+    absl::flat_hash_set<const Node*>* visiting) {
+  auto it = memo->find(node);
+  if (it != memo->end()) {
+    return it->second;
+  }
+  if (!visiting->insert(node).second) {
+    return false;
+  }
+
+  TensorShapeProto contents_proto;
+  if (TryGetContentsProtoAttr(node->attrs(), &contents_proto) &&
+      HasDynamicDimExprs(contents_proto)) {
+    return (*memo)[node] = true;
+  }
+
+  bool has_dynamic = false;
+  TensorShapeProto inferred_shape_proto;
+  if (GetNodeAttr(node->attrs(), "has_dynamic", &has_dynamic).ok() &&
+      has_dynamic &&
+      GetNodeAttr(node->attrs(), "user_inferred_shape", &inferred_shape_proto)
+          .ok() &&
+      HasDynamicDimExprs(inferred_shape_proto)) {
+    return (*memo)[node] = true;
+  }
+
+  if (GetShapeFromDirectDynamicSource(node, &inferred_shape_proto) ||
+      node->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr) {
+    return (*memo)[node] = true;
+  }
+
+  for (const Edge* edge : node->in_edges()) {
+    if (edge->IsControlEdge()) continue;
+    if (HasTransitiveDynamicShapeContents(edge->src(), memo, visiting)) {
+      return (*memo)[node] = true;
+    }
+  }
+
+  return (*memo)[node] = false;
+}
+
+bool GetConstTensor(const Node* node, Tensor* tensor) {
+  if (node == nullptr || !node->IsConstant()) {
+    return false;
+  }
+  const TensorProto* tensor_proto;
+  if (!GetNodeAttr(node->attrs(), "value", &tensor_proto).ok()) {
+    return false;
+  }
+  DataType dtype;
+  if (!GetNodeAttr(node->attrs(), "dtype", &dtype).ok()) {
+    return false;
+  }
+  *tensor = Tensor(dtype);
+  return tensor->FromProto(cpu_allocator(), *tensor_proto);
+}
+
+bool GetInputConstTensor(const Node* node, int input_index, Tensor* tensor) {
+  const Edge* edge;
+  if (!node->input_edge(input_index, &edge).ok()) {
+    return false;
+  }
+  return GetConstTensor(edge->src(), tensor);
+}
+
+bool GetTensorIntValues(const Tensor& tensor, std::vector<int64_t>* values) {
+  values->clear();
+  if (tensor.dims() == 0) {
+    values->reserve(1);
+    switch (tensor.dtype()) {
+      case DT_INT32:
+        values->push_back(tensor.scalar<int32>()());
+        return true;
+      case DT_INT64:
+        values->push_back(tensor.scalar<int64_t>()());
+        return true;
+      default:
+        return false;
+    }
+  }
+  if (tensor.dims() != 1) {
+    return false;
+  }
+  values->reserve(tensor.NumElements());
+  switch (tensor.dtype()) {
+    case DT_INT32: {
+      auto flat = tensor.flat<int32>();
+      for (int i = 0; i < flat.size(); ++i) values->push_back(flat(i));
+      return true;
+    }
+    case DT_INT64: {
+      auto flat = tensor.flat<int64_t>();
+      for (int i = 0; i < flat.size(); ++i) values->push_back(flat(i));
+      return true;
+    }
+    default:
+      return false;
+  }
+}
+
+void CopyContentAt(const TensorShapeProto& input_contents, int64_t index,
+                   TensorShapeProto* output_contents) {
+  output_contents->add_dim()->CopyFrom(input_contents.dim(index));
+  ExpressionProto* output_expression = output_contents->add_expressions();
+  if (index < input_contents.expressions_size()) {
+    output_expression->CopyFrom(input_contents.expressions(index));
+  } else {
+    output_expression->set_constant_value(input_contents.dim(index).size());
+  }
+}
+
+void AppendScalarConstantContent(int64_t value,
+                                 TensorShapeProto* output_contents) {
+  output_contents->add_dim()->set_size(value);
+  output_contents->add_expressions()->set_constant_value(value);
+}
+
+void AppendScalarContentFromTensor(const Tensor& tensor,
+                                   TensorShapeProto* output_contents) {
+  if (tensor.dtype() == DT_INT32) {
+    AppendScalarConstantContent(tensor.scalar<int32>()(), output_contents);
+  } else {
+    AppendScalarConstantContent(tensor.scalar<int64_t>()(), output_contents);
+  }
+}
+
+ExpressionProto MakeConstantExpressionProto(int64_t value) {
+  ExpressionProto expr;
+  expr.set_constant_value(value);
+  return expr;
+}
+
+ExpressionProto GetContentExpressionProto(const TensorShapeProto& contents,
+                                          int64_t index) {
+  if (index < contents.expressions_size()) {
+    return contents.expressions(index);
+  }
+  return MakeConstantExpressionProto(contents.dim(index).size());
+}
+
+ExpressionProto MakeMulExpressionProto(ExpressionProto lhs,
+                                       ExpressionProto rhs) {
+  ExpressionProto expr;
+  auto* mul = expr.mutable_mul_node();
+  *mul->mutable_lhs() = std::move(lhs);
+  *mul->mutable_rhs() = std::move(rhs);
+  return expr;
+}
+
+bool TryGetFoldedValueContents(const Node* node, int output_index,
+                               TensorShapeProto* out_contents,
+                               absl::flat_hash_set<const Node*>* visiting) {
+  out_contents->Clear();
+  if (output_index != 0) {
+    return false;
+  }
+  if (!visiting->insert(node).second) {
+    return false;
+  }
+  auto remove_from_visiting =
+      gtl::MakeCleanup([&] { visiting->erase(node); });
+
+  if (TryGetContentsProtoAttr(node->attrs(), out_contents)) {
+    return true;
+  }
+
+  bool has_dynamic = false;
+  TensorShapeProto user_inferred_shape;
+  if (GetNodeAttr(node->attrs(), "has_dynamic", &has_dynamic).ok() &&
+      has_dynamic &&
+      GetNodeAttr(node->attrs(), "user_inferred_shape", &user_inferred_shape)
+          .ok()) {
+    *out_contents = user_inferred_shape;
+    return true;
+  }
+
+  if (GetShapeFromDirectDynamicSource(node, out_contents)) {
+    return true;
+  }
+
+  auto recurse_input = [&](int input_index,
+                           TensorShapeProto* input_contents) -> bool {
+    const Edge* input_edge;
+    if (!node->input_edge(input_index, &input_edge).ok()) {
+      return false;
+    }
+    return TryGetFoldedValueContents(input_edge->src(), input_edge->src_output(),
+                                     input_contents, visiting);
+  };
+
+  if (node->IsIdentity() || node->type_string() == "Cast") {
+    return recurse_input(0, out_contents);
+  }
+
+  TensorShapeProto input_contents;
+  if (node->type_string() == "Reshape" && recurse_input(0, &input_contents)) {
+    Tensor shape_tensor;
+    std::vector<int64_t> shape_dims;
+    if (!GetInputConstTensor(node, 1, &shape_tensor) ||
+        !GetTensorIntValues(shape_tensor, &shape_dims)) {
+      return false;
+    }
+    if (input_contents.dim_size() == 1 && shape_dims.empty()) {
+      CopyContentAt(input_contents, 0, out_contents);
+      return true;
+    }
+    if (shape_dims.size() == 1 && input_contents.dim_size() == shape_dims[0]) {
+      out_contents->CopyFrom(input_contents);
+      return true;
+    }
+    return false;
+  }
+
+  if (node->type_string() == "Pack") {
+    for (int i = 0; i < node->num_inputs(); ++i) {
+      TensorShapeProto scalar_contents;
+      Tensor scalar_tensor;
+      if (recurse_input(i, &scalar_contents)) {
+        if (scalar_contents.dim_size() != 1) {
+          return false;
+        }
+        CopyContentAt(scalar_contents, 0, out_contents);
+      } else if (GetInputConstTensor(node, i, &scalar_tensor) &&
+                 scalar_tensor.dims() == 0 &&
+                 (scalar_tensor.dtype() == DT_INT32 ||
+                  scalar_tensor.dtype() == DT_INT64)) {
+        AppendScalarContentFromTensor(scalar_tensor, out_contents);
+      } else {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  if (node->type_string() == "ConcatV2") {
+    Tensor axis_tensor;
+    std::vector<int64_t> axis_values;
+    if (!GetInputConstTensor(node, node->num_inputs() - 1, &axis_tensor) ||
+        !GetTensorIntValues(axis_tensor, &axis_values) ||
+        axis_values.size() != 1) {
+      return false;
+    }
+    int64_t axis = axis_values[0];
+    if (axis != 0 && axis != -1) {
+      return false;
+    }
+    for (int i = 0; i < node->num_inputs() - 1; ++i) {
+      TensorShapeProto part_contents;
+      if (!recurse_input(i, &part_contents)) {
+        return false;
+      }
+      for (int64_t j = 0; j < part_contents.dim_size(); ++j) {
+        CopyContentAt(part_contents, j, out_contents);
+      }
+    }
+    return true;
+  }
+
+  if ((node->type_string() == "Gather" || node->type_string() == "GatherV2") &&
+      recurse_input(0, &input_contents)) {
+    Tensor indices_tensor;
+    std::vector<int64_t> indices;
+    if (!GetInputConstTensor(node, 1, &indices_tensor) ||
+        !GetTensorIntValues(indices_tensor, &indices)) {
+      return false;
+    }
+
+    int64_t axis = 0;
+    if (node->type_string() == "GatherV2") {
+      Tensor axis_tensor;
+      std::vector<int64_t> axis_values;
+      if (!GetInputConstTensor(node, 2, &axis_tensor) ||
+          !GetTensorIntValues(axis_tensor, &axis_values) ||
+          axis_values.size() != 1) {
+        return false;
+      }
+      axis = axis_values[0];
+    }
+
+    const int64_t params_rank = 1;
+    if (axis < 0) axis += params_rank;
+    if (axis != 0) {
+      return false;
+    }
+
+    const int64_t rank = input_contents.dim_size();
+    for (int64_t index : indices) {
+      if (index < 0) index += rank;
+      if (index < 0 || index >= rank) {
+        return false;
+      }
+      CopyContentAt(input_contents, index, out_contents);
+    }
+    return true;
+  }
+
+  if (node->type_string() == "Prod" && recurse_input(0, &input_contents)) {
+    Tensor reduction_indices_tensor;
+    std::vector<int64_t> axes;
+    bool keep_dims = false;
+    if (!GetInputConstTensor(node, 1, &reduction_indices_tensor) ||
+        !GetTensorIntValues(reduction_indices_tensor, &axes) ||
+        !GetNodeAttr(node->attrs(), "keep_dims", &keep_dims).ok() ||
+        keep_dims || axes.size() != 1 ||
+        (axes[0] != 0 && axes[0] != -1) || input_contents.dim_size() == 0) {
+      return false;
+    }
+    int64_t value = 1;
+    ExpressionProto expr = GetContentExpressionProto(input_contents, 0);
+    for (int64_t i = 0; i < input_contents.dim_size(); ++i) {
+      value *= input_contents.dim(i).size();
+      if (i > 0) {
+        expr = MakeMulExpressionProto(std::move(expr),
+                                      GetContentExpressionProto(input_contents, i));
+      }
+    }
+    out_contents->add_dim()->set_size(value);
+    out_contents->add_expressions()->Swap(&expr);
+    return true;
+  }
+
+  if (node->type_string() == "Slice" && recurse_input(0, &input_contents)) {
+    Tensor begin_tensor;
+    Tensor size_tensor;
+    std::vector<int64_t> begin;
+    std::vector<int64_t> size;
+    if (!GetInputConstTensor(node, 1, &begin_tensor) ||
+        !GetInputConstTensor(node, 2, &size_tensor) ||
+        !GetTensorIntValues(begin_tensor, &begin) ||
+        !GetTensorIntValues(size_tensor, &size) || begin.size() != 1 ||
+        size.size() != 1) {
+      return false;
+    }
+    int64_t start = begin[0];
+    if (start < 0 || start > input_contents.dim_size()) {
+      return false;
+    }
+    int64_t length = size[0] < 0 ? input_contents.dim_size() - start : size[0];
+    if (length < 0 || start + length > input_contents.dim_size()) {
+      return false;
+    }
+    for (int64_t i = 0; i < length; ++i) {
+      CopyContentAt(input_contents, start + i, out_contents);
+    }
+    return true;
+  }
+
+  if (node->type_string() == "StridedSlice" &&
+      recurse_input(0, &input_contents)) {
+    Tensor begin_tensor;
+    Tensor end_tensor;
+    Tensor strides_tensor;
+    std::vector<int64_t> begin;
+    std::vector<int64_t> end;
+    std::vector<int64_t> strides;
+    int64_t begin_mask = 0;
+    int64_t end_mask = 0;
+    int64_t ellipsis_mask = 0;
+    int64_t new_axis_mask = 0;
+    int64_t shrink_axis_mask = 0;
+    if (!GetInputConstTensor(node, 1, &begin_tensor) ||
+        !GetInputConstTensor(node, 2, &end_tensor) ||
+        !GetInputConstTensor(node, 3, &strides_tensor) ||
+        !GetTensorIntValues(begin_tensor, &begin) ||
+        !GetTensorIntValues(end_tensor, &end) ||
+        !GetTensorIntValues(strides_tensor, &strides) || begin.size() != 1 ||
+        end.size() != 1 || strides.size() != 1 ||
+        !GetNodeAttr(node->attrs(), "begin_mask", &begin_mask).ok() ||
+        !GetNodeAttr(node->attrs(), "end_mask", &end_mask).ok() ||
+        !GetNodeAttr(node->attrs(), "ellipsis_mask", &ellipsis_mask).ok() ||
+        !GetNodeAttr(node->attrs(), "new_axis_mask", &new_axis_mask).ok() ||
+        !GetNodeAttr(node->attrs(), "shrink_axis_mask", &shrink_axis_mask).ok()) {
+      return false;
+    }
+    if (ellipsis_mask != 0 || new_axis_mask != 0) {
+      return false;
+    }
+    const int64_t rank = input_contents.dim_size();
+    int64_t stride = strides[0];
+    if (stride == 0) {
+      return false;
+    }
+    int64_t start = (begin_mask & 1) ? (stride > 0 ? 0 : rank - 1) : begin[0];
+    int64_t stop = (end_mask & 1) ? (stride > 0 ? rank : -1) : end[0];
+    if (start < 0) start += rank;
+    if (stop < 0 && !(end_mask & 1 && stride < 0)) stop += rank;
+    if (shrink_axis_mask & 1) {
+      if (start < 0 || start >= rank) {
+        return false;
+      }
+      CopyContentAt(input_contents, start, out_contents);
+      return true;
+    }
+    if (stride < 0) {
+      return false;
+    }
+    start = std::max<int64_t>(0, start);
+    stop = std::min<int64_t>(rank, stop);
+    for (int64_t i = start; i < stop; i += stride) {
+      CopyContentAt(input_contents, i, out_contents);
+    }
+    return true;
+  }
+
+  return false;
+}
+
+bool TryGetFoldedValueContents(const Node* node, int output_index,
+                               TensorShapeProto* out_contents) {
+  absl::flat_hash_set<const Node*> visiting;
+  return TryGetFoldedValueContents(node, output_index, out_contents, &visiting);
+}
 
 // For stateless RNGs ops, they are pure but device-dependent. Those ops are not
 // constant-foldable.
@@ -242,8 +700,23 @@ bool IsConstantFoldable(
         shape_map,
     const std::function<bool(const Node*)>& consider,
     int64_t max_constant_size_in_bytes,
-    std::unordered_map<const Node*, std::vector<Tensor>>*
-        shape_replacement_map) {
+    std::unordered_map<const Node*, std::vector<Tensor>>* shape_replacement_map,
+    std::unordered_map<const Node*, bool>* dynamic_contents_memo) {
+  TensorShapeProto exact_contents;
+  const bool has_exact_contents =
+      TryGetFoldedValueContents(n, 0, &exact_contents);
+  TensorShapeProto dynamic_shape;
+  const bool has_dynamic = GetShapeFromDirectDynamicSource(n, &dynamic_shape);
+  const bool is_shape_derived =
+      n->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr;
+  absl::flat_hash_set<const Node*> dynamic_contents_visiting;
+  const bool has_transitive_dynamic_contents =
+      HasTransitiveDynamicShapeContents(n, dynamic_contents_memo,
+                                        &dynamic_contents_visiting);
+  if ((has_dynamic || is_shape_derived || has_transitive_dynamic_contents) &&
+      (!has_exact_contents || n->num_outputs() > 1)) {
+    return false;
+  }
   if (n->IsConstant()) {
     // Skip constant folding resources as they cannot be deep copied.
     return n->output_type(0) != DT_RESOURCE;
@@ -326,10 +799,11 @@ void ConsiderConstantFoldableNode(
     Node* n, const ConstantFoldingOptions& opts, std::vector<Node*>* nodes,
     std::unordered_map<const Node*, gtl::FlatSet<Node*>>* constant_control_deps,
     std::unordered_map<const Node*, std::vector<Tensor>>* shape_replacement_map,
+    std::unordered_map<const Node*, bool>* dynamic_contents_memo,
     bool* internal_node_inserted) {
   if (!IsConstantFoldable(n, opts.shape_map, opts.consider,
                           opts.max_constant_size_in_bytes,
-                          shape_replacement_map)) {
+                          shape_replacement_map, dynamic_contents_memo)) {
     return;
   }
   // A node is constant provided all of its non-control incoming Tensors come
@@ -383,16 +857,17 @@ void FindConstantFoldableNodes(
     const Graph* graph, const ConstantFoldingOptions& opts,
     std::vector<Node*>* nodes,
     std::unordered_map<const Node*, gtl::FlatSet<Node*>>* constant_control_deps,
-    std::unordered_map<const Node*, std::vector<Tensor>>*
-        shape_replacement_map) {
+    std::unordered_map<const Node*, std::vector<Tensor>>* shape_replacement_map) {
   bool internal_node_inserted = false;
+  std::unordered_map<const Node*, bool> dynamic_contents_memo;
   // Walk the nodes in data flow order.
   ReverseDFS(
       *graph, nullptr,
       [nodes, constant_control_deps, shape_replacement_map,
-       &internal_node_inserted, &opts](Node* n) {
+       &dynamic_contents_memo, &internal_node_inserted, &opts](Node* n) {
         ConsiderConstantFoldableNode(n, opts, nodes, constant_control_deps,
                                      shape_replacement_map,
+                                     &dynamic_contents_memo,
                                      &internal_node_inserted);
       },
       NodeComparatorName());
@@ -449,13 +924,29 @@ void AddShapeNodeToConstantGraph(
         shape_replacement_map,
     std::unordered_map<Node*, std::vector<Node*>>* node_map,
     const ConstantFoldNameGenerator& generate_new_name, Graph* constant_graph) {
+  TensorShapeProto user_inferred_shape;
+  const bool has_dynamic =
+      GetShapeFromDirectDynamicSource(n, &user_inferred_shape);
+  TensorShapeProto exact_contents;
+  const bool has_exact_contents =
+      TryGetFoldedValueContents(n, 0, &exact_contents);
   std::vector<Node*>& added = (*node_map)[n];
   const string& node_name = n->name();
   for (const Tensor& t : shape_replacement_map.at(n)) {
+    VLOG(1) << "Constant folding shape node " << node_name
+            << " into Const with value " << t.SummarizeValue(16);
     auto builder =
         NodeDefBuilder(generate_new_name(constant_graph, node_name), "Const")
             .Attr("dtype", t.dtype())
             .Attr("value", t);
+    if (has_dynamic) {
+      builder.Attr("has_dynamic", has_dynamic)
+          .Attr("user_inferred_shape", user_inferred_shape);
+    }
+    if (has_exact_contents && HasDynamicDimExprs(exact_contents)) {
+      builder.Attr(kUserInferredValueContentsAttrName,
+                   exact_contents.SerializeAsString());
+    }
     NodeDef def;
     CHECK(builder.Finalize(&def).ok());
     Node* constant_node;
@@ -546,6 +1037,12 @@ bool ReplaceTensorWithConstant(
                                ? DeviceType{partition_device->device_type()}
                                : DEVICE_CPU;
   if (partition_device && device_type != DEVICE_CPU) {
+    // Constant folding replaces one output edge-set at a time. Be
+    // conservative for non-CPU multi-output ops, since partially replacing a
+    // node can violate per-output placement or memory-type assumptions.
+    if (tensor.first->num_outputs() > 1) {
+      return false;
+    }
     MemoryTypeVector input_mvec;
     MemoryTypeVector output_mvec;
     if (!MemoryTypesForNode(graph->op_registry(), device_type,
@@ -574,10 +1071,40 @@ bool ReplaceTensorWithConstant(
     }
   }
   const string& node_name = n->name();
+  TensorShapeProto user_inferred_shape;
+  const bool has_dynamic =
+      GetShapeFromDirectDynamicSource(tensor.first, &user_inferred_shape);
+  const bool is_shape_derived =
+      tensor.first->attrs().FindByString(kXlaShapeDerivedAttrName) != nullptr;
+  if (tensor.second != 0 && (has_dynamic || is_shape_derived)) {
+    VLOG(1) << "Skipping replacement of " << tensor.first->name() << " :: "
+            << tensor.second
+            << " because symbolic content preservation is only supported for "
+            << "single-output replacements";
+    return false;
+  }
+  TensorShapeProto exact_contents;
+  const bool has_exact_contents =
+      TryGetFoldedValueContents(tensor.first, tensor.second, &exact_contents);
+  if ((has_dynamic || is_shape_derived) && !has_exact_contents) {
+    VLOG(1) << "Skipping replacement of " << tensor.first->name() << " :: "
+            << tensor.second
+            << " because constant folding could not preserve symbolic contents";
+    return false;
+  }
   Node* constant_node;
   auto builder = NodeDefBuilder(generate_new_name(graph, node_name), "Const")
                      .Attr("dtype", constant.dtype())
                      .Attr("value", constant);
+  if (has_dynamic) {
+    builder.Attr("has_dynamic", has_dynamic)
+        .Attr("user_inferred_shape", user_inferred_shape);
+  }
+  if (has_exact_contents && HasDynamicDimExprs(exact_contents)) {
+    builder.Attr("has_dynamic", true)
+        .Attr(kUserInferredValueContentsAttrName,
+              exact_contents.SerializeAsString());
+  }
   if (partition_device) {
     builder.Device(partition_device->name());
   }
@@ -592,6 +1119,9 @@ bool ReplaceTensorWithConstant(
 
   VLOG(1) << "Replacing " << tensor.first->name() << " :: " << tensor.second
           << " with a constant";
+  VLOG(1) << "ReplaceTensorWithConstant creating Const from "
+          << tensor.first->name() << " :: " << tensor.second
+          << " with value " << constant.SummarizeValue(16);
 
   if (!NodeBuilder(builder).Finalize(graph, &constant_node).ok()) {
     return false;
diff --git a/tensorflow/core/common_runtime/constant_folding_test.cc b/tensorflow/core/common_runtime/constant_folding_test.cc
index 481a85add4893c..fa4cf9f2b9c6a2 100644
--- a/tensorflow/core/common_runtime/constant_folding_test.cc
+++ b/tensorflow/core/common_runtime/constant_folding_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <vector>
 
 #include "tensorflow/cc/ops/array_ops_internal.h"
+#include "tensorflow/cc/ops/function_ops.h"
 #include "tensorflow/cc/ops/nn_ops.h"
 #include "tensorflow/cc/ops/sendrecv_ops.h"
 #include "tensorflow/cc/ops/standard_ops.h"
@@ -32,6 +33,7 @@ limitations under the License.
 #include "tensorflow/core/framework/node_def_util.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/graph/node_builder.h"
@@ -45,6 +47,15 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
+TensorShapeProto MakeDynamicShapeProto977x16() {
+  TensorShapeProto proto;
+  proto.add_dim()->set_size(977);
+  proto.add_dim()->set_size(16);
+  proto.add_expressions()->set_variable_id(0);
+  proto.add_expressions()->set_constant_value(16);
+  return proto;
+}
+
 class ConstantFoldingTest : public ::testing::Test {
  protected:
   template <typename T>
@@ -634,6 +645,333 @@ TEST_F(ConstantFoldingTest, ConstShapeKnown) {
   }
 }
 
+TEST_F(ConstantFoldingTest, FoldShapeFromDynamicArgPreservesContents) {
+  Graph g(OpRegistry::Global());
+  Scope s = Scope::NewRootScope();
+  auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0);
+  auto shape = ops::Shape(s.WithOpName("shape"), arg);
+  auto send = ops::_Send(s.WithOpName("send"), shape, "send", "sender", 0,
+                         "receiver");
+  TF_ASSERT_OK(s.ToGraph(&g));
+
+  std::unordered_map<string, Node*> index_by_name = g.BuildNodeNameIndex();
+  Node* arg_node = index_by_name.at("arg");
+  arg_node->AddAttr("_output_shapes",
+                    std::vector<TensorShapeProto>{MakeDynamicShapeProto977x16()});
+
+  PartialTensorShape partial_shape({977, 16});
+  std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
+  shape_map[arg_node->name()].push_back(partial_shape);
+
+  ConstantFoldingOptions opts;
+  opts.shape_map = &shape_map;
+  bool was_mutated = false;
+  TF_ASSERT_OK(
+      ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated));
+
+  index_by_name = g.BuildNodeNameIndex();
+  Node* send_node = index_by_name.at("send");
+  const Edge* send_input = nullptr;
+  TF_ASSERT_OK(send_node->input_edge(0, &send_input));
+  Node* folded = send_input->src();
+  ExpectNodeEqual<int32>(folded, {977, 16}, {2});
+
+  // This test checks the stable contract we care about: after folding
+  // Shape(arg), the replacement Const still carries symbolic contents.
+  string serialized_contents_proto;
+  TF_ASSERT_OK(GetNodeAttr(folded->attrs(),
+                           "_user_inferred_value_contents",
+                           &serialized_contents_proto));
+  TensorShapeProto contents_proto;
+  ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto));
+  ASSERT_EQ(contents_proto.expressions_size(), 2);
+  EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0)));
+  EXPECT_FALSE(IsDynamicDimExpr(contents_proto.expressions(1)));
+  EXPECT_EQ(contents_proto.dim(0).size(), 977);
+  EXPECT_EQ(contents_proto.dim(1).size(), 16);
+}
+
+TEST_F(ConstantFoldingTest, FoldSliceOfDynamicShapePreservesContents) {
+  Graph g(OpRegistry::Global());
+  Scope s = Scope::NewRootScope();
+  auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0);
+  auto shape = ops::Shape(s.WithOpName("shape"), arg);
+  auto begin = ops::Const(s.WithOpName("begin"), {0}, {1});
+  auto size = ops::Const(s.WithOpName("size"), {1}, {1});
+  auto slice = ops::Slice(s.WithOpName("slice"), shape, begin, size);
+  auto send = ops::_Send(s.WithOpName("send"), slice, "send", "sender", 0,
+                         "receiver");
+  TF_ASSERT_OK(s.ToGraph(&g));
+
+  std::unordered_map<string, Node*> index_by_name = g.BuildNodeNameIndex();
+  Node* arg_node = index_by_name.at("arg");
+  arg_node->AddAttr("_output_shapes",
+                    std::vector<TensorShapeProto>{MakeDynamicShapeProto977x16()});
+
+  PartialTensorShape partial_shape({977, 16});
+  std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
+  shape_map[arg_node->name()].push_back(partial_shape);
+
+  ConstantFoldingOptions opts;
+  opts.shape_map = &shape_map;
+  bool was_mutated = false;
+  TF_ASSERT_OK(
+      ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated));
+
+  index_by_name = g.BuildNodeNameIndex();
+  Node* send_node = index_by_name.at("send");
+  const Edge* send_input = nullptr;
+  TF_ASSERT_OK(send_node->input_edge(0, &send_input));
+  Node* folded = send_input->src();
+  ExpectNodeEqual<int32>(folded, {977}, {1});
+
+  // Folding Slice(Shape(arg), [0], [1]) should preserve the selected symbolic
+  // content, not just the concrete value 977.
+  string serialized_contents_proto;
+  TF_ASSERT_OK(GetNodeAttr(folded->attrs(),
+                           "_user_inferred_value_contents",
+                           &serialized_contents_proto));
+  TensorShapeProto contents_proto;
+  ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto));
+  ASSERT_EQ(contents_proto.expressions_size(), 1);
+  EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0)));
+  EXPECT_EQ(contents_proto.dim(0).size(), 977);
+}
+
+TEST_F(ConstantFoldingTest, FoldGatherOfDynamicShapePreservesContents) {
+  Graph g(OpRegistry::Global());
+  Scope s = Scope::NewRootScope();
+  auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0);
+  auto shape = ops::Shape(s.WithOpName("shape"), arg);
+  auto index = ops::Const(s.WithOpName("index"), 0);
+  auto gather = ops::GatherV2(s.WithOpName("gather"), shape, index,
+                              ops::Const(s.WithOpName("axis"), 0));
+  auto send = ops::_Send(s.WithOpName("send"), gather, "send", "sender", 0,
+                         "receiver");
+  TF_ASSERT_OK(s.ToGraph(&g));
+
+  std::unordered_map<string, Node*> index_by_name = g.BuildNodeNameIndex();
+  Node* arg_node = index_by_name.at("arg");
+  arg_node->AddAttr("_output_shapes",
+                    std::vector<TensorShapeProto>{MakeDynamicShapeProto977x16()});
+
+  PartialTensorShape partial_shape({977, 16});
+  std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
+  shape_map[arg_node->name()].push_back(partial_shape);
+
+  ConstantFoldingOptions opts;
+  opts.shape_map = &shape_map;
+  bool was_mutated = false;
+  TF_ASSERT_OK(
+      ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated));
+
+  index_by_name = g.BuildNodeNameIndex();
+  Node* send_node = index_by_name.at("send");
+  const Edge* send_input = nullptr;
+  TF_ASSERT_OK(send_node->input_edge(0, &send_input));
+  Node* folded = send_input->src();
+  ExpectNodeEqual<int32>(folded, {977}, {});
+
+  // Folding Gather(Shape(arg), 0, axis=0) should preserve the selected
+  // symbolic content, not just the concrete scalar value 977.
+  string serialized_contents_proto;
+  TF_ASSERT_OK(GetNodeAttr(folded->attrs(),
+                           "_user_inferred_value_contents",
+                           &serialized_contents_proto));
+  TensorShapeProto contents_proto;
+  ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto));
+  ASSERT_EQ(contents_proto.expressions_size(), 1);
+  EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0)));
+  EXPECT_EQ(contents_proto.dim(0).size(), 977);
+}
+
+TEST_F(ConstantFoldingTest, FoldReshapeOfDynamicShapePreservesContents) {
+  Graph g(OpRegistry::Global());
+  Scope s = Scope::NewRootScope();
+  auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0);
+  auto shape = ops::Shape(s.WithOpName("shape"), arg);
+  auto begin = ops::Const(s.WithOpName("begin"), {0}, {1});
+  auto size = ops::Const(s.WithOpName("size"), {1}, {1});
+  auto slice = ops::Slice(s.WithOpName("slice"), shape, begin, size);
+  auto scalar_shape = ops::Const<int32>(s.WithOpName("scalar_shape"), {});
+  auto reshape =
+      ops::Reshape(s.WithOpName("reshape"), slice, scalar_shape);
+  auto send = ops::_Send(s.WithOpName("send"), reshape, "send", "sender", 0,
+                         "receiver");
+  TF_ASSERT_OK(s.ToGraph(&g));
+
+  std::unordered_map<string, Node*> index_by_name = g.BuildNodeNameIndex();
+  Node* arg_node = index_by_name.at("arg");
+  arg_node->AddAttr("_output_shapes",
+                    std::vector<TensorShapeProto>{MakeDynamicShapeProto977x16()});
+
+  PartialTensorShape partial_shape({977, 16});
+  std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
+  shape_map[arg_node->name()].push_back(partial_shape);
+
+  ConstantFoldingOptions opts;
+  opts.shape_map = &shape_map;
+  bool was_mutated = false;
+  TF_ASSERT_OK(
+      ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated));
+
+  index_by_name = g.BuildNodeNameIndex();
+  Node* send_node = index_by_name.at("send");
+  const Edge* send_input = nullptr;
+  TF_ASSERT_OK(send_node->input_edge(0, &send_input));
+  Node* folded = send_input->src();
+  ExpectNodeEqual<int32>(folded, {977}, {});
+
+  // Folding Reshape(Slice(Shape(arg), [0], [1]), []) should preserve the
+  // selected symbolic content when the shape-vector result is scalarized.
+  string serialized_contents_proto;
+  TF_ASSERT_OK(GetNodeAttr(folded->attrs(),
+                           "_user_inferred_value_contents",
+                           &serialized_contents_proto));
+  TensorShapeProto contents_proto;
+  ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto));
+  ASSERT_EQ(contents_proto.expressions_size(), 1);
+  EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0)));
+  EXPECT_EQ(contents_proto.dim(0).size(), 977);
+}
+
+TEST_F(ConstantFoldingTest, FoldPackWithRepeatedDynamicInputPreservesContents) {
+  Graph g(OpRegistry::Global());
+  Scope s = Scope::NewRootScope();
+  auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0);
+  auto shape = ops::Shape(s.WithOpName("shape"), arg);
+  auto first = ops::GatherV2(s.WithOpName("first"), shape,
+                             ops::Const(s.WithOpName("index"), 0),
+                             ops::Const(s.WithOpName("axis"), 0));
+  // Reusing the same producer is a DAG, not a recursion cycle. Both visits must
+  // preserve its symbolic content.
+  OutputList pack_inputs = {first, first};
+  auto pack = ops::Stack(s.WithOpName("pack"), pack_inputs,
+                         ops::Stack::Axis(0));
+  auto send = ops::_Send(s.WithOpName("send"), pack, "send", "sender", 0,
+                         "receiver");
+  TF_ASSERT_OK(s.ToGraph(&g));
+
+  std::unordered_map<string, Node*> index_by_name = g.BuildNodeNameIndex();
+  Node* arg_node = index_by_name.at("arg");
+  arg_node->AddAttr("_output_shapes",
+                    std::vector<TensorShapeProto>{MakeDynamicShapeProto977x16()});
+
+  PartialTensorShape partial_shape({977, 16});
+  std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
+  shape_map[arg_node->name()].push_back(partial_shape);
+
+  ConstantFoldingOptions opts;
+  opts.shape_map = &shape_map;
+  bool was_mutated = false;
+  TF_ASSERT_OK(
+      ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated));
+
+  index_by_name = g.BuildNodeNameIndex();
+  Node* send_node = index_by_name.at("send");
+  const Edge* send_input = nullptr;
+  TF_ASSERT_OK(send_node->input_edge(0, &send_input));
+  Node* folded = send_input->src();
+  ASSERT_TRUE(folded->IsConstant());
+  ExpectNodeEqual<int32>(folded, {977, 977}, {2});
+
+  string serialized_contents_proto;
+  TF_ASSERT_OK(GetNodeAttr(folded->attrs(),
+                           "_user_inferred_value_contents",
+                           &serialized_contents_proto));
+  TensorShapeProto contents_proto;
+  ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto));
+  ASSERT_EQ(contents_proto.expressions_size(), 2);
+  EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(0)));
+  EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(1)));
+  EXPECT_EQ(contents_proto.dim(0).size(), 977);
+  EXPECT_EQ(contents_proto.dim(1).size(), 977);
+}
+
+TEST_F(ConstantFoldingTest, FoldPackKeepsSymbolicContentsAligned) {
+  Graph g(OpRegistry::Global());
+  Scope s = Scope::NewRootScope();
+  auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0);
+  auto shape = ops::Shape(s.WithOpName("shape"), arg);
+  auto dynamic_value = ops::GatherV2(
+      s.WithOpName("dynamic_value"), shape,
+      ops::Const(s.WithOpName("index"), 0),
+      ops::Const(s.WithOpName("axis"), 0));
+  auto static_value = ops::Const(s.WithOpName("static_value"), 16);
+  auto pack = ops::Stack(s.WithOpName("pack"),
+                         OutputList{static_value, dynamic_value},
+                         ops::Stack::Axis(0));
+  auto send = ops::_Send(s.WithOpName("send"), pack, "send", "sender", 0,
+                         "receiver");
+  TF_ASSERT_OK(s.ToGraph(&g));
+
+  std::unordered_map<string, Node*> index_by_name = g.BuildNodeNameIndex();
+  Node* arg_node = index_by_name.at("arg");
+  arg_node->AddAttr("_output_shapes",
+                    std::vector<TensorShapeProto>{MakeDynamicShapeProto977x16()});
+
+  PartialTensorShape partial_shape({977, 16});
+  std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
+  shape_map[arg_node->name()].push_back(partial_shape);
+
+  ConstantFoldingOptions opts;
+  opts.shape_map = &shape_map;
+  bool was_mutated = false;
+  TF_ASSERT_OK(
+      ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated));
+
+  index_by_name = g.BuildNodeNameIndex();
+  const Edge* send_input = nullptr;
+  TF_ASSERT_OK(index_by_name.at("send")->input_edge(0, &send_input));
+  Node* folded = send_input->src();
+  ASSERT_TRUE(folded->IsConstant());
+  ExpectNodeEqual<int32>(folded, {16, 977}, {2});
+
+  string serialized_contents_proto;
+  TF_ASSERT_OK(GetNodeAttr(folded->attrs(),
+                           "_user_inferred_value_contents",
+                           &serialized_contents_proto));
+  TensorShapeProto contents_proto;
+  ASSERT_TRUE(contents_proto.ParseFromString(serialized_contents_proto));
+  ASSERT_EQ(contents_proto.expressions_size(), 2);
+  EXPECT_FALSE(IsDynamicDimExpr(contents_proto.expressions(0)));
+  EXPECT_TRUE(IsDynamicDimExpr(contents_proto.expressions(1)));
+  EXPECT_EQ(contents_proto.dim(0).size(), 16);
+  EXPECT_EQ(contents_proto.dim(1).size(), 977);
+}
+
+TEST_F(ConstantFoldingTest, DoNotFoldUnsupportedDynamicContentsTransform) {
+  Graph g(OpRegistry::Global());
+  Scope s = Scope::NewRootScope();
+  auto arg = ops::_Arg(s.WithOpName("arg"), DT_FLOAT, 0);
+  auto shape = ops::Shape(s.WithOpName("shape"), arg);
+  auto added = ops::Add(s.WithOpName("added"), shape, shape);
+  auto send = ops::_Send(s.WithOpName("send"), added, "send", "sender", 0,
+                         "receiver");
+  TF_ASSERT_OK(s.ToGraph(&g));
+
+  std::unordered_map<string, Node*> index_by_name = g.BuildNodeNameIndex();
+  Node* arg_node = index_by_name.at("arg");
+  arg_node->AddAttr("_output_shapes",
+                    std::vector<TensorShapeProto>{MakeDynamicShapeProto977x16()});
+
+  PartialTensorShape partial_shape({977, 16});
+  std::unordered_map<string, std::vector<PartialTensorShape>> shape_map;
+  shape_map[arg_node->name()].push_back(partial_shape);
+
+  ConstantFoldingOptions opts;
+  opts.shape_map = &shape_map;
+  bool was_mutated = false;
+  TF_ASSERT_OK(
+      ConstantFold(opts, nullptr, Env::Default(), nullptr, &g, &was_mutated));
+
+  index_by_name = g.BuildNodeNameIndex();
+  Node* send_node = index_by_name.at("send");
+  const Edge* send_input = nullptr;
+  TF_ASSERT_OK(send_node->input_edge(0, &send_input));
+  EXPECT_FALSE(send_input->src()->IsConstant());
+}
+
 TEST_F(ConstantFoldingTest, NoReplacePartialOutput) {
   Graph g(OpRegistry::Global());
   {
diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 09142e303e3e13..79f4bd855fd10b 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -62,6 +62,7 @@ package(
 exports_files(
     srcs = [
         "allocator_registry.h",
+        "batch_size_resource.h",
         "collective.h",
         "control_flow.h",
         "dataset.h",
@@ -195,6 +196,7 @@ filegroup(
         "allocator.h",
         "allocator_registry.h",
         "attr_value_util.h",
+        "batch_size_resource.h",
         "bfloat16.h",
         "bounds_check.h",
         "cancellation.h",
@@ -340,6 +342,8 @@ filegroup(
         "tensor_key.h",
         "tensor_shape.cc",
         "tensor_shape.h",
+        "tensor_shape_expr.cc",
+        "tensor_shape_expr.h",
         "tensor_types.h",
         "tensor_util.h",
         "tracking_allocator.h",
@@ -702,6 +706,7 @@ cc_library(
     ],
     deps = [
         ":bounds_check",
+        ":tensor_shape_expr",
         ":tensor_shape_proto_cc",
         ":types_proto_cc",
         "//tensorflow/core/lib/core:errors",
@@ -718,10 +723,43 @@ cc_library(
         "//tensorflow/core/platform:statusor",
         "//tensorflow/core/util:overflow",
         "@eigen_archive//:eigen3",
+        "@local_xla//xla:shape_util",
     ],
     alwayslink = 1,
 )
 
+cc_library(
+    name = "tensor_shape_expr",
+    srcs = ["tensor_shape_expr.cc"],
+    hdrs = ["tensor_shape_expr.h"],
+    visibility = [
+        "//tensorflow/core:__pkg__",
+        "//tensorflow/core/grappler:__subpackages__",
+        "//tensorflow/core/runtime_fallback:__subpackages__",
+        "//tensorflow/core/tfrt/utils:__subpackages__",
+    ],
+    deps = [
+        ":tensor_shape_proto_cc",
+        "//tensorflow/core:lib",
+        "@local_xla//xla:parse_flags_from_env",
+        "@local_xla//xla:shape_util",
+        "@local_xla//xla/tsl/util:command_line_flags",
+    ],
+    alwayslink = 1,
+)
+
+tf_cc_test(
+    name = "tensor_shape_expr_test",
+    size = "small",
+    srcs = ["tensor_shape_expr_test.cc"],
+    deps = [
+        ":tensor_shape",
+        ":tensor_shape_expr",
+        "//tensorflow/core/platform:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
 cc_library(
     name = "resource_base",
     hdrs = ["resource_base.h"],
@@ -921,6 +959,7 @@ cc_library(
         ":node_def_util",
         ":op_def_proto_cc",
         ":tensor_shape",
+        ":tensor_shape_expr",
         ":tensor_shape_proto_cc",
         "//tensorflow/core/lib/core:errors",
         "//tensorflow/core/lib/core:status",
diff --git a/tensorflow/core/framework/batch_size_resource.h b/tensorflow/core/framework/batch_size_resource.h
new file mode 100644
index 00000000000000..45bdfad02699bb
--- /dev/null
+++ b/tensorflow/core/framework/batch_size_resource.h
@@ -0,0 +1,40 @@
+/* Copyright 2026 The TensorFlow Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_FRAMEWORK_BATCH_SIZE_RESOURCE_H_
+#define TENSORFLOW_CORE_FRAMEWORK_BATCH_SIZE_RESOURCE_H_
+
+#include <cstddef>
+
+#include "tensorflow/core/framework/resource_mgr.h"
+
+namespace tensorflow {
+const string BatchSizeResourceName = "BatchSizeResource_";
+class BatchSizeResource : public ResourceBase {
+ public:
+  ~BatchSizeResource() override {
+    VLOG(1) << "BatchSizeResource destroyed with batch size: " << batch_size_;
+  }
+  string DebugString() const override { return BatchSizeResourceName; }
+  void SetBatchSize(size_t s) { batch_size_ = s; }
+  size_t GetBatchSize() { return batch_size_; }
+
+ private:
+  size_t batch_size_ = 0;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_FRAMEWORK_BATCH_SIZE_RESOURCE_H_
diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc
index aefa2a416310e2..42cf77245b8e1b 100644
--- a/tensorflow/core/framework/common_shape_fns.cc
+++ b/tensorflow/core/framework/common_shape_fns.cc
@@ -2369,7 +2369,13 @@ absl::Status BroadcastBinaryOpOutputShapeFnHelper(InferenceContext* c,
       } else if (dim_y.SameHandle(dim_x)) {
         dims.push_back(dim_x);
       } else if (!c->ValueKnown(dim_x) && !c->ValueKnown(dim_y)) {
-        dims.push_back(c->UnknownDim());
+        DimensionHandle merged;
+        absl::Status s = c->Merge(dim_x, dim_y, &merged);
+        if (s.ok()) {
+          dims.push_back(merged);
+        } else {
+          dims.push_back(c->UnknownDim());
+        }
       } else {
         if (!incompatible_shape_error) {
           *out = c->UnknownShape();
diff --git a/tensorflow/core/framework/shape_inference.cc b/tensorflow/core/framework/shape_inference.cc
index b63269f68c3368..dbd0ef831871dd 100644
--- a/tensorflow/core/framework/shape_inference.cc
+++ b/tensorflow/core/framework/shape_inference.cc
@@ -249,6 +249,16 @@ void InferenceContext::ShapeHandleToProto(ShapeHandle handle,
     } else {
       dim_shape->set_size(-1);
     }
+
+    ExpressionProto* expression = proto->add_expressions();
+    if (DimExpr* expr = GetDimExpr(dim)) {
+      DimExprToProto(*expr, expression);
+    } else if (ValueKnown(dim)) {
+      expression->set_constant_value(Value(dim));
+    } else {
+      DimExprToProto(DimExpr::Unknown(xla::kMissingExpressionSentinel),
+                     expression);
+    }
   }
 }
 
@@ -282,6 +292,37 @@ DimensionHandle InferenceContext::NumElements(ShapeHandle s) {
   }
 }
 
+DimensionHandle InferenceContext::UnknownDimWithExpr(
+    std::unique_ptr<DimExpr> expr) {
+  DimExpr* owned = shape_manager_.OwnExpr(std::move(expr));
+  return shape_manager_.MakeDim(kUnknownDim, owned);
+}
+
+DimExpr* InferenceContext::GetDimExpr(DimensionHandle d) const {
+  if (!d.IsSet()) return nullptr;
+  return d->expr_;
+}
+
+DimExpr* InferenceContext::MakeConstExpr(int64_t v) {
+  return shape_manager_.OwnExpr(
+      std::make_unique<DimExpr>(DimExpr::Const(v)));
+}
+
+DimExpr* InferenceContext::ExprForDim(DimensionHandle d) {
+  if (!d.IsSet()) return nullptr;
+
+  // If already tagged with expr, use it.
+  if (DimExpr* e = GetDimExpr(d)) return e;
+
+  // Known dim -> const expr.
+  if (ValueKnown(d)) {
+    return MakeConstExpr(Value(d));
+  }
+
+  // Unknown dim with no expr -> cannot form expression.
+  return nullptr;
+}
+
 string InferenceContext::DebugString(ShapeHandle s) {
   if (RankKnown(s)) {
     std::vector<string> vals;
@@ -923,8 +964,40 @@ absl::Status InferenceContext::MakeShapeFromShapeProto(
     const TensorShapeProto& proto, ShapeHandle* out) {
   *out = nullptr;
   TF_RETURN_IF_ERROR(PartialTensorShape::IsValidShape(proto));
-  PartialTensorShape partial_shape(proto);
-  return MakeShapeFromPartialTensorShape(partial_shape, out);
+
+  if (proto.unknown_rank()) {
+    *out = UnknownShape();
+    return absl::OkStatus();
+  }
+
+  std::vector<DimensionHandle> dims;
+  dims.reserve(proto.dim_size());
+  for (int i = 0; i < proto.dim_size(); ++i) {
+    const auto& dim_proto = proto.dim(i);
+    if (dim_proto.size() >= 0) {
+      // Known dimension
+      dims.push_back(MakeDim(dim_proto.size()));
+    } else {
+      // Unknown dimension - check for expression.
+      if (i < proto.expressions_size() &&
+          proto.expressions(i).node_type_case() !=
+              ExpressionProto::NODE_TYPE_NOT_SET) {
+        // Deserialize expression
+        DimExpr expr = DimExprFromProto(proto.expressions(i));
+        if (expr) {
+          DimExpr* owned = shape_manager_.OwnExpr(
+              std::make_unique<DimExpr>(std::move(expr)));
+          dims.push_back(shape_manager_.MakeDim(kUnknownDim, owned));
+        } else {
+          dims.push_back(UnknownDim());
+        }
+      } else {
+        dims.push_back(UnknownDim());
+      }
+    }
+  }
+  *out = MakeShape(dims);
+  return absl::OkStatus();
 }
 
 absl::Status InferenceContext::GetScalarFromTensor(const Tensor* t,
@@ -1030,24 +1103,40 @@ absl::Status InferenceContext::Divide(DimensionHandle dividend,
                                       DimensionOrConstant divisor,
                                       bool evenly_divisible,
                                       DimensionHandle* out) {
-  const int64_t divisor_value = Value(divisor);
-  if (divisor_value == 1) {
+  const bool dividend_known = ValueKnown(dividend);
+  const bool divisor_known = ValueKnown(divisor);
+
+  // Validate divisor if known.
+  if (divisor_known && Value(divisor) <= 0) {
+    return errors::InvalidArgument("Divisor must be positive but is ",
+                                   Value(divisor));
+  }
+  // Fast-path: x / 1 = x
+  if (divisor_known && Value(divisor) == 1) {
     *out = dividend;
-  } else if (!ValueKnown(dividend) ||
-             (divisor.dim.IsSet() && !ValueKnown(divisor.dim))) {
-    *out = UnknownDim();
-  } else {
+    return absl::OkStatus();
+  }
+  // If both known, do numeric divide.
+  if (dividend_known && divisor_known) {
     const int64_t v = Value(dividend);
-    if (divisor_value <= 0) {
-      return errors::InvalidArgument("Divisor must be positive but is ",
-                                     divisor_value);
-    }
-    if (evenly_divisible && (v % divisor_value) != 0) {
+    const int64_t d = Value(divisor);
+    if (evenly_divisible && (v % d) != 0) {
       return errors::InvalidArgument(
-          "Dimension size must be evenly divisible by ", divisor_value,
-          " but is ", v);
+          "Dimension size must be evenly divisible by ", d, " but is ", v);
     }
-    *out = MakeDim(v / divisor_value);
+    *out = MakeDim(v / d);
+    return absl::OkStatus();
+  }
+  // At least one operand unknown: try to build expression.
+  DimExpr* lhs = ExprForDim(dividend);
+  DimExpr* rhs = divisor.dim.IsSet() ? ExprForDim(divisor.dim)
+                                     : MakeConstExpr(divisor.val);
+  if (lhs && rhs) {
+    DimExpr* node =
+        shape_manager_.OwnExpr(std::make_unique<DimExpr>(*lhs / *rhs));
+    *out = shape_manager_.MakeDim(kUnknownDim, node);
+  } else {
+    *out = UnknownDim();  // Can't form expr.
   }
   return absl::OkStatus();
 }
@@ -1055,26 +1144,42 @@ absl::Status InferenceContext::Divide(DimensionHandle dividend,
 absl::Status InferenceContext::Add(DimensionHandle first,
                                    DimensionOrConstant second,
                                    DimensionHandle* out) {
-  const int64_t first_value = Value(first);
-  const int64_t second_value = Value(second);
-  // Special cases.
-  if (first_value == 0) {
+  const bool first_known = ValueKnown(first);
+  const bool second_known = ValueKnown(second);
+
+  // Fast-path: x + 0 = x
+  if (first_known && Value(first) == 0) {
     *out = MakeDim(second);
-  } else if (second_value == 0) {
+    return absl::OkStatus();
+  }
+  if (second_known && Value(second) == 0) {
     *out = first;
-  } else if (first_value == kUnknownDim || second_value == kUnknownDim) {
-    *out = UnknownDim();
-  } else {
-    // Invariant: Both values are known and positive. Still in run-time we can
-    // get pair of values which cannot be store in output. Check below will
-    // report error. We still need to avoid undefined behavior of signed
-    // overflow and use unsigned addition.
-    const int64_t sum = static_cast<uint64>(first_value) + second_value;
+    return absl::OkStatus();
+  }
+
+  // If both known, do numeric add.
+  if (first_known && second_known) {
+    const int64_t sum = static_cast<uint64_t>(Value(first)) +
+                        static_cast<uint64_t>(Value(second));
     if (sum < 0) {
       return errors::InvalidArgument("Dimension size overflow from adding ",
-                                     first_value, " and ", second_value);
+                                     Value(first), " and ", Value(second));
     }
     *out = MakeDim(sum);
+    return absl::OkStatus();
+  }
+
+  // At least one operand unknown: try to build expression.
+  DimExpr* lhs = ExprForDim(first);
+  DimExpr* rhs =
+      second.dim.IsSet() ? ExprForDim(second.dim) : MakeConstExpr(second.val);
+
+  if (lhs && rhs) {
+    DimExpr* node =
+        shape_manager_.OwnExpr(std::make_unique<DimExpr>(*lhs + *rhs));
+    *out = shape_manager_.MakeDim(kUnknownDim, node);
+  } else {
+    *out = UnknownDim();  // Can't form expr.
   }
   return absl::OkStatus();
 }
@@ -1082,22 +1187,35 @@ absl::Status InferenceContext::Add(DimensionHandle first,
 absl::Status InferenceContext::Subtract(DimensionHandle first,
                                         DimensionOrConstant second,
                                         DimensionHandle* out) {
-  const int64_t first_value = Value(first);
-  const int64_t second_value = Value(second);
-  // Special cases.
-  if (second_value == 0) {
+  const bool first_known = ValueKnown(first);
+  const bool second_known = ValueKnown(second);
+  // Fast-path: x - 0 = x
+  if (second_known && Value(second) == 0) {
     *out = first;
-  } else if (first_value == kUnknownDim || second_value == kUnknownDim) {
-    *out = UnknownDim();
-  } else {
-    // Invariant: Both values are known, first_value is non-negative, and
-    // second_value is positive.
+    return absl::OkStatus();
+  }
+  // If both known, do numeric subtract.
+  if (first_known && second_known) {
+    const int64_t first_value = Value(first);
+    const int64_t second_value = Value(second);
     if (first_value < second_value) {
       return errors::InvalidArgument(
           "Negative dimension size caused by subtracting ", second_value,
           " from ", first_value);
     }
     *out = MakeDim(first_value - second_value);
+    return absl::OkStatus();
+  }
+  // At least one operand unknown: try to build expression.
+  DimExpr* lhs = ExprForDim(first);
+  DimExpr* rhs =
+      second.dim.IsSet() ? ExprForDim(second.dim) : MakeConstExpr(second.val);
+  if (lhs && rhs) {
+    DimExpr* node =
+        shape_manager_.OwnExpr(std::make_unique<DimExpr>(*lhs - *rhs));
+    *out = shape_manager_.MakeDim(kUnknownDim, node);
+  } else {
+    *out = UnknownDim();  // Can't form expr.
   }
   return absl::OkStatus();
 }
@@ -1105,21 +1223,31 @@ absl::Status InferenceContext::Subtract(DimensionHandle first,
 absl::Status InferenceContext::Multiply(DimensionHandle first,
                                         DimensionOrConstant second,
                                         DimensionHandle* out) {
+  const bool first_known = ValueKnown(first);
+  const bool second_known = ValueKnown(second);
   const int64_t first_value = Value(first);
   const int64_t second_value = Value(second);
-  // Special cases.
-  if (first_value == 0) {
+
+  // Fast-paths for identity and zero cases.
+  if (first_known && first_value == 0) {
     *out = first;
-  } else if (second_value == 0) {
+    return absl::OkStatus();
+  }
+  if (second_known && second_value == 0) {
     *out = MakeDim(second);
-  } else if (first_value == 1) {
+    return absl::OkStatus();
+  }
+  if (first_known && first_value == 1) {
     *out = MakeDim(second);
-  } else if (second_value == 1) {
+    return absl::OkStatus();
+  }
+  if (second_known && second_value == 1) {
     *out = first;
-  } else if (first_value == kUnknownDim || second_value == kUnknownDim) {
-    *out = UnknownDim();
-  } else {
-    // Invariant: Both values are known and greater than 1.
+    return absl::OkStatus();
+  }
+
+  // If both known, do numeric multiply.
+  if (first_known && second_known) {
     const int64_t product = MultiplyWithoutOverflow(first_value, second_value);
     if (product < 0) {
       return errors::InvalidArgument(
@@ -1127,6 +1255,20 @@ absl::Status InferenceContext::Multiply(DimensionHandle first,
           first_value, " and ", second_value);
     }
     *out = MakeDim(product);
+    return absl::OkStatus();
+  }
+
+  // At least one operand unknown: try to build expression.
+  DimExpr* lhs = ExprForDim(first);
+  DimExpr* rhs =
+      second.dim.IsSet() ? ExprForDim(second.dim) : MakeConstExpr(second.val);
+
+  if (lhs && rhs) {
+    DimExpr* node =
+        shape_manager_.OwnExpr(std::make_unique<DimExpr>(*lhs * *rhs));
+    *out = shape_manager_.MakeDim(kUnknownDim, node);
+  } else {
+    *out = UnknownDim();  // Can't form expr.
   }
   return absl::OkStatus();
 }
diff --git a/tensorflow/core/framework/shape_inference.h b/tensorflow/core/framework/shape_inference.h
index 8bfd301d860de1..d1670e348c62f7 100644
--- a/tensorflow/core/framework/shape_inference.h
+++ b/tensorflow/core/framework/shape_inference.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include "absl/memory/memory.h"
 #include "tensorflow/core/framework/full_type.pb.h"
 #include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/macros.h"
@@ -116,13 +117,15 @@ class InferenceContext;
 class Dimension {
  private:
   Dimension();
-  Dimension(int64_t value);
+  Dimension(int64_t value, DimExpr* expr = nullptr);
   ~Dimension() {}
 
   const int64_t value_;
+  DimExpr* expr_;
 
   friend class InferenceContext;
   friend class ShapeManager;
+  friend class ::tensorflow::grappler::SymbolicShapeManager;
   Dimension(const Dimension&) = delete;
   void operator=(const Dimension&) = delete;
 };
@@ -578,6 +581,20 @@ class InferenceContext {
 
   inline DimensionHandle UnknownDim() { return MakeDim(kUnknownDim); }
 
+  // Create a new unknown dimension (size = -1) tagged with a DimExpr.
+  // The expression is owned by this context's ShapeManager.
+  DimensionHandle UnknownDimWithExpr(std::unique_ptr<DimExpr> expr);
+  // Return the expression pointer for a dimension, or nullptr if none.
+  DimExpr* GetDimExpr(DimensionHandle d) const;
+  // Creates a constant DimExpr node for the given value.
+  // The expression is owned by this context's ShapeManager.
+  DimExpr* MakeConstExpr(int64_t v);
+  // Returns the Expr representation for the given dimension:
+  // - If dim has an expr, returns it
+  // - If dim is known, returns a new Const expr
+  // - If dim is unknown with no expr, returns nullptr
+  DimExpr* ExprForDim(DimensionHandle d);
+
   // Returns in <val> a scalar value from an input tensor <t>.  The input tensor
   // must be a 0-dimensional int32 or int64 tensor.  Caller must ensure that the
   // input tensor is not NULL.
@@ -743,8 +760,6 @@ class InferenceContext {
 
   // Adds new outputs; useful when mutating the graph.
   absl::Status ExpandOutputs(int new_output_size);
-
- private:
   // Creates and stores shapes for use in InferenceContext.
   class ShapeManager {
    public:
@@ -760,21 +775,32 @@ class InferenceContext {
 
     // Returns a new dimension of the given size.  The returned value
     // is owned by this class.
-    inline DimensionHandle MakeDim(DimensionOrConstant d) {
+    inline DimensionHandle MakeDim(DimensionOrConstant d,
+                                   DimExpr* expr = nullptr) {
       if (d.dim.IsSet()) {
         return d.dim;
       } else {
-        all_dims_.push_back(new Dimension(d.val));
+        all_dims_.push_back(new Dimension(d.val, expr));
         return all_dims_.back();
       }
     }
+    // Takes ownership of an expression and returns a raw pointer to it.
+    DimExpr* OwnExpr(std::unique_ptr<DimExpr> expr) {
+      if (!expr) return nullptr;
+      DimExpr* ptr = expr.get();
+      all_exprs_.push_back(std::move(expr));
+      return ptr;
+    }
 
    private:
     std::vector<Shape*> all_shapes_;    // values are owned.
     std::vector<Dimension*> all_dims_;  // values are owned.
+    std::vector<std::unique_ptr<DimExpr>> all_exprs_;   // expressions are owned.
   };
+ private:
 
   friend class ::tensorflow::grappler::GraphProperties;
+  friend class ::tensorflow::grappler::SymbolicShapeManager;
 
   friend class ShapeInferenceTest;      // For testing Relax functions.
   friend class ShapeInferenceTestutil;  // For testing shapes.
@@ -888,8 +914,10 @@ class InferenceContext {
 // -----------------------------------------------------------------------------
 // Template and inline method implementations, please ignore
 
-inline Dimension::Dimension() : value_(InferenceContext::kUnknownDim) {}
-inline Dimension::Dimension(int64_t value) : value_(value) {
+inline Dimension::Dimension()
+    : value_(InferenceContext::kUnknownDim), expr_(nullptr) {}
+inline Dimension::Dimension(int64_t value, DimExpr* expr)
+    : value_(value), expr_(expr) {
   DCHECK(value >= 0 || value == InferenceContext::kUnknownDim)
       << "Dimension must be non-negative or equal to "
          "InferenceContext::kUnknownDim but got "
diff --git a/tensorflow/core/framework/shape_inference_test.cc b/tensorflow/core/framework/shape_inference_test.cc
index c5dbc299b86540..9a10a05ed65163 100644
--- a/tensorflow/core/framework/shape_inference_test.cc
+++ b/tensorflow/core/framework/shape_inference_test.cc
@@ -1035,6 +1035,35 @@ TEST_F(ShapeInferenceTest, KnownShapeToProto) {
   EXPECT_FALSE(proto.unknown_rank());
   EXPECT_EQ(3, proto.dim_size());
   EXPECT_EQ(1, proto.dim(0).size());
+  ASSERT_EQ(3, proto.expressions_size());
+  EXPECT_EQ(1, proto.expressions(0).constant_value());
+  EXPECT_EQ(2, proto.expressions(1).constant_value());
+  EXPECT_EQ(3, proto.expressions(2).constant_value());
+  EXPECT_FALSE(proto.dim(0).has_expr());
+}
+
+TEST_F(ShapeInferenceTest, DynamicExpressionShapeProtoRoundTrip) {
+  NodeDef def;
+  std::vector<ShapeHandle> empty;
+  InferenceContext c(kVersion, def, MakeOpDef(0, 2), empty, {}, {}, {});
+
+  TensorShapeProto input_proto;
+  input_proto.add_dim()->set_size(-1);
+  DimExprToProto(DimExpr::Var(7), input_proto.add_expressions());
+
+  ShapeHandle shape;
+  TF_ASSERT_OK(c.MakeShapeFromShapeProto(input_proto, &shape));
+  TensorShapeProto proto;
+  c.ShapeHandleToProto(shape, &proto);
+
+  ASSERT_EQ(1, proto.expressions_size());
+  EXPECT_EQ(DimExpr::Var(7), DimExprFromProto(proto.expressions(0)));
+  EXPECT_FALSE(proto.dim(0).has_expr());
+
+  ShapeHandle restored;
+  TF_ASSERT_OK(c.MakeShapeFromShapeProto(proto, &restored));
+  ASSERT_NE(nullptr, c.GetDimExpr(c.Dim(restored, 0)));
+  EXPECT_EQ(DimExpr::Var(7), *c.GetDimExpr(c.Dim(restored, 0)));
 }
 
 TEST_F(ShapeInferenceTest, UnknownShapeToProto) {
diff --git a/tensorflow/core/framework/tensor_shape.cc b/tensorflow/core/framework/tensor_shape.cc
index 35c628216ed3c6..4b06e0861d86cf 100644
--- a/tensorflow/core/framework/tensor_shape.cc
+++ b/tensorflow/core/framework/tensor_shape.cc
@@ -17,15 +17,24 @@ limitations under the License.
 
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/util/overflow.h"
+#include "xla/printer.h"
 
 namespace tensorflow {
 
+std::string ExprToString(const xla::DExpr& e) {
+  if (!e && !e.is_unknown()) return "";
+  xla::StringPrinter printer;
+  e->print(&printer);
+  return std::move(printer).ToString();
+}
+
 // TensorShape and PartialTensorShape should have no fields beyond
 // TensorShapeRep.  In particular, their sizes should be the same.
 static_assert(sizeof(TensorShapeRep) == sizeof(TensorShape),
@@ -152,6 +161,11 @@ TensorShapeBase<Shape>::TensorShapeBase(const TensorShapeProto& proto) {
     for (const auto& d : proto.dim()) {
       AddDim(d.size());
     }
+    if (TensorShapeExpressionsEnabled()) {
+      for (const auto& e : proto.expressions()) {
+        AddExpression(DimExprFromProto(e));
+      }
+    }
   }
 }
 
@@ -191,6 +205,11 @@ absl::Status TensorShapeBase<Shape>::BuildTensorShapeBase(
         }
       }
     }
+    if (TensorShapeExpressionsEnabled()) {
+      for (const auto& e : proto.expressions()) {
+        out->AddExpression(DimExprFromProto(e));
+      }
+    }
   }
   return absl::OkStatus();
 }
@@ -375,10 +394,46 @@ void TensorShapeRep::Clear() {
   set_data_type(DT_INVALID);
 }
 
+void TensorShapeRep::set_expression(int d, xla::DExpr expr) {
+  if (!TensorShapeExpressionsEnabled()) {
+    expressions_.clear();
+    return;
+  }
+  if (expressions_.size() <= static_cast<size_t>(d)) {
+    expressions_.resize(d + 1,
+                        xla::DExpr::Unknown(xla::kMissingExpressionSentinel));
+  }
+  expressions_[d] = expr ? std::move(expr)
+                         : xla::DExpr::Unknown(xla::kMissingExpressionSentinel);
+}
+
+void TensorShapeRep::AddExpression(xla::DExpr expr) {
+  if (!TensorShapeExpressionsEnabled()) {
+    return;
+  }
+  CHECK_LT(expressions_.size(), ndims_byte());
+  expressions_.push_back(expr ? std::move(expr)
+                              : xla::DExpr::Unknown(
+                                    xla::kMissingExpressionSentinel));
+}
+
+void TensorShapeRep::set_expressions(std::vector<xla::DExpr> exprs) {
+  if (!TensorShapeExpressionsEnabled()) {
+    expressions_.clear();
+    return;
+  }
+  CHECK_LE(exprs.size(), ndims_byte());
+  for (auto& expr : exprs) {
+    if (!expr) expr = xla::DExpr::Unknown(xla::kMissingExpressionSentinel);
+  }
+  expressions_ = std::move(exprs);
+}
+
 void TensorShapeRep::ClearAllButDataType() {
   if (tag() == REP_OUT_OF_LINE) {
     delete as64()->dims_;
   }
+  expressions_.clear();
   set_tag(REP16);
   set_ndims_byte(0);
   // Leaves data_type alone
@@ -505,6 +560,9 @@ void TensorShapeBase<Shape>::UnsafeAddDim(int64_t size,
 template <class Shape>
 void TensorShapeBase<Shape>::AppendShape(const TensorShapeBase& shape) {
   for (auto d : shape) AddDim(d.size);
+  for (auto e : shape.get_expressions()){
+     AddExpression(e);
+  }
 }
 
 template <class Shape>
@@ -585,6 +643,13 @@ template <class Shape>
 void TensorShapeBase<Shape>::set_dim(int d, int64_t size) {
   CHECK_GE(d, 0);
   CHECK_LT(d, dims());
+  // After DExpr migration, missing slots may be normalized to Unknown().
+  // Preserve those placeholders here instead of materializing them into a
+  // concrete constant just because the dimension size changed.
+  if (get_expressions().size() > d &&
+      get_expression(d).kind() != xla::DExpr::Kind::kUnknown) {
+    set_expression(d, xla::DExpr::Const(size));
+  }
   if (!kIsPartial) {
     CHECK_GE(size, 0);
   }
@@ -646,6 +711,13 @@ absl::Status TensorShapeBase<Shape>::SetDimWithStatus(int d, int64_t size) {
     }
   }
 
+  // After DExpr migration, missing slots may be normalized to Unknown().
+  // Preserve those placeholders here instead of materializing them into a
+  // concrete constant just because the dimension size changed.
+  if (get_expressions().size() > d &&
+      get_expression(d).kind() != xla::DExpr::Kind::kUnknown) {
+    set_expression(d, xla::DExpr::Const(size));
+  }
   return RecomputeNumElements();
 }
 
@@ -661,11 +733,31 @@ void TensorShapeBase<Shape>::RemoveDimRange(int begin, int end) {
   if (begin >= end) return;
   absl::InlinedVector<int64_t, 8UL> vals;
   AppendTo(*this, &vals);
+  std::vector<xla::DExpr> new_exprs(get_expressions().begin(),
+                                    get_expressions().end());
+  if (begin < static_cast<int64_t>(new_exprs.size())) {
+    int64_t expr_end = end;
+    if (expr_end > static_cast<int64_t>(new_exprs.size())) {
+      expr_end = new_exprs.size();
+    }
+    if (expr_end > begin) {
+      new_exprs.erase(new_exprs.begin() + begin, new_exprs.begin() + expr_end);
+    }
+  }
+
   vals.erase(vals.begin() + begin, vals.begin() + end);
+
+  // Truncate if the removed dims reduce rank below expression vector size.
+  const int64_t new_rank = vals.size();
+  if (new_exprs.size() > static_cast<size_t>(new_rank)) {
+    new_exprs.resize(new_rank);
+  }
+
   ClearAllButDataType();
   for (auto dval : vals) {
     AddDim(dval);
   }
+  set_expressions(new_exprs);
   TF_CHECK_OK(RecomputeNumElements());
 }
 
@@ -700,7 +792,26 @@ absl::Status TensorShapeBase<Shape>::RemoveDimRangeWithStatus(int begin,
 
   absl::InlinedVector<int64_t, 8UL> vals;
   AppendTo(*this, &vals);
+
+  std::vector<xla::DExpr> new_exprs(get_expressions().begin(),
+                                    get_expressions().end());
+  if (begin < static_cast<int64_t>(new_exprs.size())) {
+    int64_t expr_end = end;
+    if (expr_end > static_cast<int64_t>(new_exprs.size())) {
+      expr_end = new_exprs.size();
+    }
+    if (expr_end > begin) {
+      new_exprs.erase(new_exprs.begin() + begin, new_exprs.begin() + expr_end);
+    }
+  }
+
   vals.erase(vals.begin() + begin, vals.begin() + end);
+
+  const int64_t new_rank = vals.size();
+  if (new_exprs.size() > static_cast<size_t>(new_rank)) {
+    new_exprs.resize(new_rank);
+  }
+
   ClearAllButDataType();
 
   absl::Status s = absl::OkStatus();
@@ -710,7 +821,7 @@ absl::Status TensorShapeBase<Shape>::RemoveDimRangeWithStatus(int begin,
       return s;
     }
   }
-
+  set_expressions(new_exprs);
   return RecomputeNumElements();
 }
 
@@ -731,6 +842,12 @@ void TensorShapeBase<Shape>::AsProto(TensorShapeProto* proto) const {
     for (int i = 0; i < dims(); i++) {
       proto->add_dim()->set_size(dim_size(i));
     }
+    if (TensorShapeExpressionsEnabled()) {
+      for (int i = 0; i < get_expressions().size(); i++) {
+        ExpressionProto* eproto = proto->add_expressions();
+        DimExprToProto(get_expression(i), eproto);
+      }
+    }
   }
 }
 
@@ -764,6 +881,11 @@ string TensorShapeRep::DebugString() const {
     } else {
       strings::StrAppend(&s, dim);
     }
+    if (shape.get_expression(i)) {
+      strings::StrAppend(&s, "<");
+      strings::StrAppend(&s, ExprToString(shape.get_expression(i)));
+      strings::StrAppend(&s, ">");
+    }
   }
   strings::StrAppend(&s, "]");
   return s;
@@ -787,6 +909,16 @@ string TensorShapeRep::DebugString(const TensorShapeProto& proto) {
     first = false;
   }
   strings::StrAppend(&s, "]");
+  if (TensorShapeExpressionsEnabled()) {
+    strings::StrAppend(&s, "<");
+    first = true;
+    for (const auto& e : proto.expressions()) {
+      if (!first) strings::StrAppend(&s, ",");
+      strings::StrAppend(&s, ExprToString(DimExprFromProto(e)));
+      first = false;
+    }
+    strings::StrAppend(&s, ">");
+  }
   return s;
 }
 
@@ -950,6 +1082,8 @@ absl::Status PartialTensorShape::MergeWith(const PartialTensorShape& shape,
       return s;
     }
   }
+  result->set_expressions(std::vector<xla::DExpr>(shape.get_expressions().begin(),
+                                                  shape.get_expressions().end()));
   return absl::OkStatus();
 }
 
diff --git a/tensorflow/core/framework/tensor_shape.h b/tensorflow/core/framework/tensor_shape.h
index 0bcf1fc54af844..9a9e5159a6556f 100644
--- a/tensorflow/core/framework/tensor_shape.h
+++ b/tensorflow/core/framework/tensor_shape.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_H_
 
 #include <string>
+#include <vector>
 
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/types.pb.h"
@@ -28,6 +29,7 @@ limitations under the License.
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/statusor.h"
+#include "xla/shape_expr.h"
 
 namespace tensorflow {
 
@@ -73,7 +75,63 @@ class TensorShapeRep {
   std::string DebugString() const;
   static std::string DebugString(const TensorShapeProto& proto);
 
+  void set_expression(int d, xla::DExpr expr);
+
+  void AddExpression(xla::DExpr expr);
+
+  // Set the array of dynamic multipliers.
+  void set_expressions(std::vector<xla::DExpr> exprs);
+
+  // Get the array of dynamic multipliers.
+  absl::Span<const xla::DExpr> get_expressions() const {
+    return expressions_;
+  }
+
+  // Get the array of dynamic multipliers, filling missing entries with
+  // constant expressions derived from the concrete dimensions.
+  std::vector<xla::DExpr> get_filled_expressions() const {
+    if (ndims_byte() == kUnknownRank) {
+      return {};
+    }
+    std::vector<xla::DExpr> exprs(ndims_byte());
+    for (int i = 0; i < ndims_byte(); ++i) {
+      exprs[i] = get_filled_expression(i);
+    }
+    return exprs;
+  }
+
+  // Return the multiplier for a specific dynamic dimension.
+  // -1 if the dimension is not dynamic.
+  const xla::DExpr& get_expression(int64_t dimension) const {
+    static const xla::DExpr kMissingExpression =
+        xla::DExpr::Unknown(xla::kMissingExpressionSentinel);
+    if (dimension < 0) return kMissingExpression;
+    const size_t dim = static_cast<size_t>(dimension);
+    if (dim >= expressions_.size()) {
+      return kMissingExpression;
+    }
+    return expressions_[dim];
+  }
+
+  xla::DExpr get_filled_expression(int64_t dimension) const {
+    if (dimension < 0) {
+      return xla::DExpr::Unknown(xla::kMissingExpressionSentinel);
+    }
+    const size_t dim = static_cast<size_t>(dimension);
+    if (dim < expressions_.size() && expressions_[dim]) {
+      return expressions_[dim];
+    }
+
+    if (ndims_byte() == kUnknownRank || dim >= ndims_byte()) {
+      return xla::DExpr::Unknown(xla::kMissingExpressionSentinel);
+    }
+
+    return constant_expression_for_dim(dim);
+  }
+
  protected:
+  std::vector<xla::DExpr> expressions_;
+
   // Constructable only via TensorShapeBase
   TensorShapeRep() = default;
 
@@ -143,6 +201,20 @@ class TensorShapeRep {
   void set_num_elements(int64_t n) { num_elements_ = n; }
 
  private:
+  xla::DExpr constant_expression_for_dim(size_t dim) const {
+    int64_t dim_value = -1;
+    if (tag() == REP16) {
+      uint16 raw_dim = as16()->dims_[dim];
+      dim_value = raw_dim == kUnknownRep16 ? -1 : raw_dim;
+    } else if (tag() == REP32) {
+      uint32 raw_dim = as32()->dims_[dim];
+      dim_value = raw_dim == kUnknownRep32 ? -1 : raw_dim;
+    } else {
+      dim_value = (*as64()->dims_)[dim];
+    }
+    return xla::DExpr::Const(dim_value);
+  }
+
   void DestructorOutOfLine();
   void SlowCopyFrom(const TensorShapeRep& b);
 
@@ -710,6 +782,7 @@ absl::Status TensorShape::AsEigenDSizesWithPaddingWithStatus(
 
 inline TensorShapeRep::TensorShapeRep(const TensorShapeRep& b) {
   num_elements_ = b.num_elements_;
+  expressions_ = b.expressions_;
   if (b.tag() != REP_OUT_OF_LINE) {
     memcpy(buf(), b.buf(), sizeof(u_.buf));
     // memcpy above Implicitly does:
@@ -723,6 +796,7 @@ inline TensorShapeRep::TensorShapeRep(const TensorShapeRep& b) {
 
 inline TensorShapeRep::TensorShapeRep(TensorShapeRep&& b) {
   num_elements_ = b.num_elements_;
+  expressions_ = b.expressions_;
   memcpy(buf(), b.buf(), sizeof(u_.buf));
   // memcpy above Implicitly does:
   //   set_ndims_byte(b.ndims_byte());
@@ -738,6 +812,8 @@ inline TensorShapeRep::~TensorShapeRep() {
 
 inline void TensorShapeRep::operator=(const TensorShapeRep& b) {
   num_elements_ = b.num_elements_;
+  expressions_ = b.expressions_;
+
   if (tag() != REP_OUT_OF_LINE && b.tag() != REP_OUT_OF_LINE) {
     memcpy(buf(), b.buf(), sizeof(u_.buf));
     // memcpy above implicitly also does:
@@ -753,6 +829,8 @@ inline void TensorShapeRep::operator=(TensorShapeRep&& b) {
     DestructorOutOfLine();
   }
   num_elements_ = b.num_elements_;
+  expressions_ = b.expressions_;
+
   memcpy(buf(), b.buf(), sizeof(u_.buf));
   // memcpy above Implicitly does:
   //   set_ndims_byte(b.ndims_byte());
diff --git a/tensorflow/core/framework/tensor_shape.proto b/tensorflow/core/framework/tensor_shape.proto
index 45d5b78ecbbc4c..efddff41dc64ba 100644
--- a/tensorflow/core/framework/tensor_shape.proto
+++ b/tensorflow/core/framework/tensor_shape.proto
@@ -22,6 +22,10 @@ message TensorShapeProto {
 
     // Optional name of the tensor dimension.
     string name = 2;
+    //Only keep one symbolic expr.
+    //Symbolic expression for this dimension when size == -1.
+    //Allows tracking relationships between unknown dimensions.
+    ExpressionProto expr = 3;
   };
 
   // Dimensions of the tensor, such as {"input", 30}, {"output", 40}
@@ -43,4 +47,57 @@ message TensorShapeProto {
   //
   // If true, "dim.size()" must be 0.
   bool unknown_rank = 3;
+
+  repeated ExpressionProto expressions = 4;
+
 };
+
+message ExpressionProto {
+  oneof node_type {
+    int32 constant_value = 1;  // cons
+    int32 variable_id = 2;     // var
+    AddNode add_node = 3;      // exp + exp
+    SubNode sub_node = 4;      // exp - exp
+    MulNode mul_node = 5;      // exp * exp
+    DivNode div_node = 6;      // exp / exp
+    MaxNode max_node = 7;      // max(exp, exp)
+    GtNode gt_node = 8;        // exp > exp
+    SelectNode select_node = 9;  // select(pred, on_true, on_false)
+  }
+}
+
+message AddNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message SubNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message MulNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message DivNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message MaxNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message GtNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message SelectNode {
+  ExpressionProto pred = 1;
+  ExpressionProto on_true = 2;
+  ExpressionProto on_false = 3;
+}
diff --git a/tensorflow/core/framework/tensor_shape_expr.cc b/tensorflow/core/framework/tensor_shape_expr.cc
new file mode 100644
index 00000000000000..8406c550477c1b
--- /dev/null
+++ b/tensorflow/core/framework/tensor_shape_expr.cc
@@ -0,0 +1,214 @@
+/* Copyright 2026 The TensorFlow Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/tensor_shape_expr.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "xla/parse_flags_from_env.h"
+#include "xla/printer.h"
+#include "xla/tsl/util/command_line_flags.h"
+
+namespace tensorflow {
+
+namespace {
+
+bool ParseTensorShapeExpressionsEnabled() {
+  bool tf_xla_enable_dynamic_sizes = false;
+  std::vector<tsl::Flag> flag_list = {
+      tsl::Flag("tf_xla_enable_dynamic_sizes", &tf_xla_enable_dynamic_sizes,
+                "XLA flag for enabling XLA dynamic sizes."),
+  };
+  xla::ParseFlagsFromEnvAndIgnoreUnknown("TF_XLA_FLAGS", flag_list);
+  return tf_xla_enable_dynamic_sizes;
+}
+
+std::optional<bool>& TensorShapeExpressionsEnabledOverride() {
+  static auto* enabled_override = new std::optional<bool>();
+  return *enabled_override;
+}
+
+void DynExprToTensorFlowProto(const xla::DynExpr& expr,
+                              ExpressionProto* proto) {
+  proto->Clear();
+  switch (expr.kind()) {
+    case xla::DExpr::Kind::kUnknown:
+      return;
+    case xla::DExpr::Kind::kConstant:
+      proto->set_constant_value(
+          static_cast<const xla::Constant&>(expr).get_val());
+      return;
+    case xla::DExpr::Kind::kVariable:
+      proto->set_variable_id(
+          static_cast<const xla::Variable&>(expr).get_id());
+      return;
+    case xla::DExpr::Kind::kAdd: {
+      const auto& add = static_cast<const xla::Add&>(expr);
+      DynExprToTensorFlowProto(*add.get_lhs(),
+                               proto->mutable_add_node()->mutable_lhs());
+      DynExprToTensorFlowProto(*add.get_rhs(),
+                               proto->mutable_add_node()->mutable_rhs());
+      return;
+    }
+    case xla::DExpr::Kind::kSub: {
+      const auto& sub = static_cast<const xla::Sub&>(expr);
+      DynExprToTensorFlowProto(*sub.get_lhs(),
+                               proto->mutable_sub_node()->mutable_lhs());
+      DynExprToTensorFlowProto(*sub.get_rhs(),
+                               proto->mutable_sub_node()->mutable_rhs());
+      return;
+    }
+    case xla::DExpr::Kind::kMul: {
+      const auto& mul = static_cast<const xla::Mul&>(expr);
+      DynExprToTensorFlowProto(*mul.get_lhs(),
+                               proto->mutable_mul_node()->mutable_lhs());
+      DynExprToTensorFlowProto(*mul.get_rhs(),
+                               proto->mutable_mul_node()->mutable_rhs());
+      return;
+    }
+    case xla::DExpr::Kind::kDiv: {
+      const auto& div = static_cast<const xla::Div&>(expr);
+      DynExprToTensorFlowProto(*div.get_lhs(),
+                               proto->mutable_div_node()->mutable_lhs());
+      DynExprToTensorFlowProto(*div.get_rhs(),
+                               proto->mutable_div_node()->mutable_rhs());
+      return;
+    }
+    case xla::DExpr::Kind::kMax: {
+      const auto& max = static_cast<const xla::MaxExpr&>(expr);
+      DynExprToTensorFlowProto(*max.get_lhs(),
+                               proto->mutable_max_node()->mutable_lhs());
+      DynExprToTensorFlowProto(*max.get_rhs(),
+                               proto->mutable_max_node()->mutable_rhs());
+      return;
+    }
+    case xla::DExpr::Kind::kGt: {
+      const auto& gt = static_cast<const xla::GtExpr&>(expr);
+      DynExprToTensorFlowProto(*gt.get_lhs(),
+                               proto->mutable_gt_node()->mutable_lhs());
+      DynExprToTensorFlowProto(*gt.get_rhs(),
+                               proto->mutable_gt_node()->mutable_rhs());
+      return;
+    }
+    case xla::DExpr::Kind::kSelect: {
+      const auto& select = static_cast<const xla::SelectExpr&>(expr);
+      DynExprToTensorFlowProto(
+          *select.get_pred(), proto->mutable_select_node()->mutable_pred());
+      DynExprToTensorFlowProto(
+          *select.get_on_true(),
+          proto->mutable_select_node()->mutable_on_true());
+      DynExprToTensorFlowProto(
+          *select.get_on_false(),
+          proto->mutable_select_node()->mutable_on_false());
+      return;
+    }
+  }
+}
+
+}  // namespace
+
+bool TensorShapeExpressionsEnabled() {
+  if (TensorShapeExpressionsEnabledOverride().has_value()) {
+    return *TensorShapeExpressionsEnabledOverride();
+  }
+  static const bool enabled = ParseTensorShapeExpressionsEnabled();
+  return enabled;
+}
+
+void SetTensorShapeExpressionsEnabledForTesting(std::optional<bool> enabled) {
+  TensorShapeExpressionsEnabledOverride() = enabled;
+}
+
+DimExpr DimExprFromProto(const ExpressionProto& proto) {
+  switch (proto.node_type_case()) {
+    case ExpressionProto::kConstantValue:
+      return DimExpr::Const(proto.constant_value());
+    case ExpressionProto::kVariableId:
+      return DimExpr::Var(proto.variable_id());
+    case ExpressionProto::kAddNode: {
+      return DimExprFromProto(proto.add_node().lhs()) +
+             DimExprFromProto(proto.add_node().rhs());
+    }
+    case ExpressionProto::kSubNode: {
+      return DimExprFromProto(proto.sub_node().lhs()) -
+             DimExprFromProto(proto.sub_node().rhs());
+    }
+    case ExpressionProto::kMulNode: {
+      return DimExprFromProto(proto.mul_node().lhs()) *
+             DimExprFromProto(proto.mul_node().rhs());
+    }
+    case ExpressionProto::kDivNode: {
+      return DimExprFromProto(proto.div_node().lhs()) /
+             DimExprFromProto(proto.div_node().rhs());
+    }
+    case ExpressionProto::kMaxNode: {
+      return DimExpr::Max(DimExprFromProto(proto.max_node().lhs()),
+                          DimExprFromProto(proto.max_node().rhs()));
+    }
+    case ExpressionProto::kGtNode: {
+      return DimExpr::Gt(DimExprFromProto(proto.gt_node().lhs()),
+                         DimExprFromProto(proto.gt_node().rhs()));
+    }
+    case ExpressionProto::kSelectNode: {
+      return DimExpr::Select(
+          DimExprFromProto(proto.select_node().pred()),
+          DimExprFromProto(proto.select_node().on_true()),
+          DimExprFromProto(proto.select_node().on_false()));
+    }
+    case ExpressionProto::NODE_TYPE_NOT_SET:
+    default:
+      return DimExpr::Unknown(xla::kMissingExpressionSentinel);
+  }
+}
+
+void DimExprToProto(const DimExpr& expr, ExpressionProto* proto) {
+  if (!expr) {
+    proto->Clear();
+    return;
+  }
+  DynExprToTensorFlowProto(*expr, proto);
+}
+
+std::string DimExprDebugString(const DimExpr& expr) {
+  if (!expr) return "_";
+  xla::StringPrinter printer;
+  expr->print(&printer);
+  return std::move(printer).ToString();
+}
+
+DimExpr* SimplifyExpr(DimExpr* expr,
+                      std::vector<std::unique_ptr<DimExpr>>* arena) {
+  if (expr == nullptr) return nullptr;
+  auto owned = std::make_unique<DimExpr>(expr->simplify());
+  DimExpr* result = owned.get();
+  arena->push_back(std::move(owned));
+  return result;
+}
+
+bool IsDynamicDimExpr(const ExpressionProto& proto) {
+  DimExpr expr = DimExprFromProto(proto);
+  return expr && expr->is_dynamic();
+}
+
+bool HasDynamicDimExprs(const TensorShapeProto& proto) {
+  for (const auto& expr : proto.expressions()) {
+    if (IsDynamicDimExpr(expr)) return true;
+  }
+  return false;
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/framework/tensor_shape_expr.h b/tensorflow/core/framework/tensor_shape_expr.h
new file mode 100644
index 00000000000000..7c61d0d079eb25
--- /dev/null
+++ b/tensorflow/core/framework/tensor_shape_expr.h
@@ -0,0 +1,57 @@
+/* Copyright 2026 The TensorFlow Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_EXPR_H_
+#define TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_EXPR_H_
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "xla/shape_expr.h"
+
+namespace tensorflow {
+
+// TensorFlow shape inference and XLA use the same owning expression value.
+// TensorFlow-specific helpers below only bridge TensorShapeProto's protobuf.
+using DimExpr = xla::DExpr;
+
+DimExpr DimExprFromProto(const ExpressionProto& proto);
+void DimExprToProto(const DimExpr& expr, ExpressionProto* proto);
+std::string DimExprDebugString(const DimExpr& expr);
+
+// Simplifies through xla::DExpr and stores the returned value in `arena`.
+DimExpr* SimplifyExpr(DimExpr* expr,
+                      std::vector<std::unique_ptr<DimExpr>>* arena);
+
+// Returns whether TensorShape should preserve symbolic expressions. The
+// shape-expression support follows the `tf_xla_enable_dynamic_sizes` flag.
+bool TensorShapeExpressionsEnabled();
+
+// Overrides TensorShapeExpressionsEnabled for tests. Passing std::nullopt
+// restores the default environment-derived behavior.
+void SetTensorShapeExpressionsEnabledForTesting(std::optional<bool> enabled);
+
+// Returns true if the expression proto depends on a symbolic variable.
+bool IsDynamicDimExpr(const ExpressionProto& proto);
+
+// Returns true if any expression attached to the TensorShapeProto is dynamic.
+bool HasDynamicDimExprs(const TensorShapeProto& proto);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_FRAMEWORK_TENSOR_SHAPE_EXPR_H_
diff --git a/tensorflow/core/framework/tensor_shape_expr_test.cc b/tensorflow/core/framework/tensor_shape_expr_test.cc
new file mode 100644
index 00000000000000..55e0f1d138f768
--- /dev/null
+++ b/tensorflow/core/framework/tensor_shape_expr_test.cc
@@ -0,0 +1,105 @@
+/* Copyright 2026 The TensorFlow Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/tensor_shape_expr.h"
+
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+class TensorShapeExpressionsEnabledTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    SetTensorShapeExpressionsEnabledForTesting(true);
+  }
+
+  void TearDown() override {
+    SetTensorShapeExpressionsEnabledForTesting(std::nullopt);
+  }
+};
+
+TEST(TensorShapeExprTest, UsesXlaCanonicalization) {
+  ExpressionProto proto;
+  auto* add = proto.mutable_add_node();
+  add->mutable_lhs()->mutable_div_node()->mutable_lhs()->set_variable_id(1);
+  add->mutable_lhs()->mutable_div_node()->mutable_rhs()->set_constant_value(2);
+  add->mutable_rhs()->mutable_div_node()->mutable_lhs()->set_variable_id(1);
+  add->mutable_rhs()->mutable_div_node()->mutable_rhs()->set_constant_value(2);
+
+  auto expr = DimExprFromProto(proto);
+  EXPECT_EQ(DimExprDebugString(expr.simplify()), "A");
+}
+
+TEST(TensorShapeExprTest, TensorFlowProtoRoundTripsSharedExpression) {
+  DimExpr original = (DimExpr::Var(7) + DimExpr::Const(3)) / 2;
+  ExpressionProto proto;
+  DimExprToProto(original, &proto);
+
+  auto round_tripped = DimExprFromProto(proto);
+  EXPECT_TRUE(xla::DynExpr::equal(original.get(), round_tripped.get()));
+}
+
+TEST(TensorShapeExprTest, TensorFlowProtoRoundTripsConditionalExpression) {
+  DimExpr variable = DimExpr::Var(7);
+  DimExpr original = DimExpr::Select(
+      DimExpr::Gt(variable, DimExpr::Const(3)),
+      DimExpr::Max(variable, DimExpr::Const(8)),
+      (variable + DimExpr::Const(3)) / 2);
+  ExpressionProto proto;
+  DimExprToProto(original, &proto);
+
+  auto round_tripped = DimExprFromProto(proto);
+  EXPECT_TRUE(xla::DynExpr::equal(original.get(), round_tripped.get()));
+}
+
+TEST(TensorShapeExprTest, SimplifyExprUsesSharedImplementation) {
+  DimExpr original =
+      (DimExpr::Var(1) / 2) + (DimExpr::Var(1) / 2);
+  std::vector<std::unique_ptr<DimExpr>> arena;
+
+  DimExpr* simplified = SimplifyExpr(&original, &arena);
+
+  ASSERT_NE(simplified, nullptr);
+  EXPECT_EQ(DimExprDebugString(*simplified), "A");
+}
+
+TEST_F(TensorShapeExpressionsEnabledTest,
+       SetExpressionsRejectsEntriesBeyondRank) {
+  TensorShape shape({2, 3});
+  EXPECT_DEATH(
+      shape.set_expressions(
+          {xla::DExpr::Var(1), xla::DExpr::Var(2), xla::DExpr::Var(3)}),
+      "");
+}
+
+TEST_F(TensorShapeExpressionsEnabledTest,
+       RemoveDimRangePreservesRemainingExpressions) {
+  TensorShape shape({2, 3});
+  shape.set_expressions({xla::DExpr::Var(1), xla::DExpr::Var(2)});
+
+  shape.RemoveDim(0);
+
+  ASSERT_EQ(shape.get_expressions().size(), 1);
+  EXPECT_TRUE(shape.get_expression(0) == xla::DExpr::Var(2));
+}
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/graph/subgraph.cc b/tensorflow/core/graph/subgraph.cc
index bb47f37ef7fbe3..d1616e2396515e 100644
--- a/tensorflow/core/graph/subgraph.cc
+++ b/tensorflow/core/graph/subgraph.cc
@@ -79,6 +79,20 @@ absl::Status FeedInputs(
     TF_RETURN_IF_ERROR(
         feed_rewrites[i]->AddNode(g, {n, id.second}, &feed_node));
 
+    // Set an attribute in _Arg node to indicate it has a batch dimension
+    auto node_attrs = n->attrs();
+    const AttrValue* shape_attr = node_attrs.FindByString("_output_shapes");
+    if (shape_attr && shape_attr->has_list()) {
+      const TensorShapeProto& shape = shape_attr->list().shape(0);
+      for (int i = 0; i < shape.dim_size(); ++i) {
+        if (shape.dim(i).size() == -1) {
+          feed_node->AddAttr("_dynamic_dim", i);
+          break;
+        }
+      }
+      // Keep _output_shapes for further runs of shape inference
+      feed_node->AddAttr("_output_shapes", *shape_attr);
+    }
     // Update name_index
     (*name_index)[feed_node->name()] = feed_node;
     // Duplicate control edges aren't allowed, but feed_node was *just* created
diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
index 613b12bb18ae3a..339ba1c7ed9353 100644
--- a/tensorflow/core/grappler/costs/graph_properties.cc
+++ b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/framework/versions.pb.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/graph/tensor_id.h"
 #include "tensorflow/core/grappler/costs/utils.h"
 #include "tensorflow/core/grappler/mutable_graph_view.h"
@@ -38,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/gtl/flatset.h"
 #include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/util/env_var.h"
 
 namespace tensorflow {
 namespace grappler {
@@ -54,11 +56,31 @@ using TensorVector = absl::InlinedVector<TensorValue, 4UL>;
 // Some ops treat "-1" specially, different from UnknownDim:
 // e.g., shape input to Reshape op.
 const int64_t kUnknownDimFromConst = INT64_MAX;
+constexpr char kInputArgSharedSymbolIdEnvVar[] =
+    "TF_GRAPPLER_INPUT_ARG_SHARED_SYMBOL_ID";
 
 // Skip const value instantiation if the number of elements in a const tensor
 // is greater than this threshold.
 const int kThresholdToSkipConstTensorInstantiation = 128;
 
+bool IsPlaceholderOp(absl::string_view op_name) {
+  return op_name == "Placeholder" || op_name == "PlaceholderV2" ||
+         op_name == "PlaceholderWithDefault";
+}
+
+int LoadInputArgSharedSymbolId() {
+  int64_t symbol_id = -1;
+  absl::Status status =
+      ReadInt64FromEnvVar(kInputArgSharedSymbolIdEnvVar, -1, &symbol_id);
+  if (!status.ok()) {
+    VLOG(1) << "Failed to read " << kInputArgSharedSymbolIdEnvVar
+            << ": " << status
+            << ". Falling back to default shared symbol id -1.";
+    return -1;
+  }
+  return static_cast<int>(symbol_id);
+}
+
 template <typename Handle>
 struct HashHandle {
   std::size_t operator()(const Handle& h) const {
@@ -189,6 +211,9 @@ class DisjointSet {
 
   absl::Status Merge(Handle x, Handle y);
   const typename HandleToObject<Handle>::Object GetMergedValue(Handle value);
+  // Returns a pointer that uniquely identifies the set containing `value`.
+  // This can be used as a stable key for associating metadata with a set.
+  void* RootId(Handle value) { return static_cast<void*>(Find(value)); }
 
  private:
   // All the handles that belong to the same set are part of the same tree, and
@@ -648,17 +673,29 @@ class SymbolicShapeRefiner {
   explicit SymbolicShapeRefiner(
       const GraphView& graph,
       const absl::flat_hash_map<string, absl::flat_hash_set<int>>& fed_ports,
-      const bool aggressive_shape_inference)
+      const bool aggressive_shape_inference,
+      const bool enable_dynamic_value_inference,
+      absl::flat_hash_set<string> top_level_input_placeholders)
       : graph_(graph),
         function_library_(OpRegistry::Global(), graph.graph()->library()),
         fed_ports_(fed_ports),
-        aggressive_shape_inference_(aggressive_shape_inference) {
+        top_level_input_placeholders_(
+            std::move(top_level_input_placeholders)),
+        input_arg_shared_symbol_id_(LoadInputArgSharedSymbolId()),
+        aggressive_shape_inference_(aggressive_shape_inference),
+        enable_dynamic_value_inference_(enable_dynamic_value_inference) {
     graph_def_version_ = graph.graph()->versions().producer();
     node_to_context_.reserve(graph.graph()->node_size());
   }
 
   const GraphView& graph() const { return graph_; }
 
+  bool IsInputArg(const NodeDef& node) const {
+    return node.op() == "_Arg" &&
+           top_level_input_placeholders_.find(node.name()) !=
+               top_level_input_placeholders_.end();
+  }
+
   struct NodeContext {
     const OpRegistrationData* op_data;
     DataTypeVector input_types;
@@ -941,7 +978,9 @@ class SymbolicShapeRefiner {
     TF_RETURN_IF_ERROR(gp.InferStatically(
         /*assume_valid_feeds=*/true,
         /*aggressive_shape_inference=*/aggressive_shape_inference_,
-        /*include_tensor_values=*/true));
+        /*include_input_tensor_values=*/true,
+        /*include_output_tensor_values=*/true,
+        /*enable_dynamic_value_inference=*/enable_dynamic_value_inference_));
 
     // Add return nodes for output shapes.
     int output = 0;
@@ -1148,32 +1187,32 @@ class SymbolicShapeRefiner {
   }
 
   struct ShapeId {
-    const NodeDef* node;
+    std::string node_name;
     int port_id;
 
     friend bool operator==(const ShapeId& lhs, const ShapeId& rhs) {
-      return lhs.node == rhs.node && lhs.port_id == rhs.port_id;
+      return lhs.node_name == rhs.node_name && lhs.port_id == rhs.port_id;
     }
 
     template <typename H>
     friend H AbslHashValue(H h, const ShapeId& s) {
-      return H::combine(std::move(h), s.node, s.port_id);
+      return H::combine(std::move(h), s.node_name, s.port_id);
     }
   };
 
   struct DimId {
-    const NodeDef* node;
+    std::string node_name;
     int port_id;
     int dim_index;
 
     friend bool operator==(const DimId& lhs, const DimId& rhs) {
-      return lhs.node == rhs.node && lhs.port_id == rhs.port_id &&
+      return lhs.node_name == rhs.node_name && lhs.port_id == rhs.port_id &&
              lhs.dim_index == rhs.dim_index;
     }
 
     template <typename H>
     friend H AbslHashValue(H h, const DimId& d) {
-      return H::combine(std::move(h), d.node, d.port_id, d.dim_index);
+      return H::combine(std::move(h), d.node_name, d.port_id, d.dim_index);
     }
   };
 
@@ -1391,7 +1430,7 @@ class SymbolicShapeRefiner {
   // Return the one ShapeHandle used to denote a fully unknown shape for a node
   // output.
   ShapeHandle GetUnknownOutputShape(const NodeDef* node, int index) {
-    ShapeId id{node, index};
+    ShapeId id{node->name(), index};
     auto it = unknown_shapes_.find(id);
     if (it != unknown_shapes_.end()) {
       return it->second;
@@ -1405,16 +1444,39 @@ class SymbolicShapeRefiner {
   // node output.
   DimensionHandle GetUnknownOutputDim(const NodeDef* node, int index,
                                       int dim_id) {
-    DimId id{node, index, dim_id};
+    DimId id{node->name(), index, dim_id};
     auto it = unknown_dims_.find(id);
     if (it != unknown_dims_.end()) {
       return it->second;
     }
     InferenceContext* c = GetContext(node);
-    DimensionHandle dim = c->UnknownDim();
+    int var_id =
+        IsInputArg(*node) ? input_arg_shared_symbol_id_
+                          : GetOrCreateStableVarId(id);
+
+    DimensionHandle dim = c->UnknownDimWithExpr(
+        std::make_unique<DimExpr>(DimExpr::Var(var_id)));
+
+    VLOG(1) << "[EXPR] GetUnknownOutputDim: node = " << node->name()
+            << " op = " << node->op()
+            << " out = " << index
+            << " dim = " << dim_id
+            << " -> Var(" << var_id << ")";
+
+    // Create an unknown dim with Var(var_id) expression.
     unknown_dims_[id] = dim;
     return dim;
   }
+  // Get or create a stable integer variable ID for a given DimId.
+  int GetOrCreateStableVarId(const DimId& id) {
+    auto it = stable_var_ids_.find(id);
+    if (it != stable_var_ids_.end()) {
+      return it->second;
+    }
+    int var_id = next_var_id_++;
+    stable_var_ids_[id] = var_id;
+    return var_id;
+  }
 
   // Returns true if all the output tensors have known values.
   bool AllOutputValuesKnown(NodeContext* c) {
@@ -1712,7 +1774,62 @@ class SymbolicShapeRefiner {
     // but instantiate a new UnknownDim to prevent incorrect symbolic shape
     // inference through UnknownDim from Const.
     InferenceContext* ic = c->inference_context.get();
+    const std::string& op = node.op();
+    const bool is_bin =
+        (op == "Sub" || op == "Add" || op == "Mul" || op == "Div");
     if (!is_fed) {
+      // Fall through to regular value inference when symbolic shape-value
+      // propagation does not apply.
+      if (enable_dynamic_value_inference_ && is_bin &&
+          c->input_tensors_as_shapes_to_propagate.size() >= 2) {
+        auto va = c->input_tensors_as_shapes_to_propagate[0];
+        auto vb = c->input_tensors_as_shapes_to_propagate[1];
+
+        if (!va.SameHandle(tensorflow::shape_inference::ShapeHandle()) &&
+            !vb.SameHandle(tensorflow::shape_inference::ShapeHandle()) &&
+            ic->RankKnown(va) && ic->RankKnown(vb) &&
+            ic->Rank(va) == ic->Rank(vb)) {
+          std::vector<tensorflow::shape_inference::DimensionHandle> out_elems;
+          out_elems.reserve(ic->Rank(va));
+          const auto is_unknown_from_const = [&](DimensionHandle dim) {
+            return ic->ValueKnown(dim) &&
+                   ic->Value(dim) == kUnknownDimFromConst;
+          };
+          bool symbolic_propagation_succeeded = true;
+
+          for (int i = 0; i < ic->Rank(va); ++i) {
+            auto da = ic->Dim(va, i);
+            auto db = ic->Dim(vb, i);
+            if (is_unknown_from_const(da) || is_unknown_from_const(db)) {
+              symbolic_propagation_succeeded = false;
+              break;
+            }
+
+            tensorflow::shape_inference::DimensionHandle r;
+            absl::Status status;
+            if (op == "Sub")
+              status = ic->Subtract(da, db, &r);
+            else if (op == "Add")
+              status = ic->Add(da, db, &r);
+            else if (op == "Mul")
+              status = ic->Multiply(da, db, &r);
+            else
+              status =
+                  ic->Divide(da, db, /*evenly_divisible=*/false, &r);
+            if (!status.ok()) {
+              symbolic_propagation_succeeded = false;
+              break;
+            }
+            out_elems.push_back(r);
+          }
+          if (symbolic_propagation_succeeded) {
+            c->output_tensors_as_shapes.resize(1);
+            c->output_tensors_as_shapes[0] = ic->MakeShape(out_elems);
+            return absl::OkStatus();
+          }
+        }
+      }
+
       if (IsConstant(node)) {
         const TensorProto& tensor_proto = node.attr().at("value").tensor();
         c->output_tensor_protos.resize(1);
@@ -1720,6 +1837,18 @@ class SymbolicShapeRefiner {
         c->output_tensors_as_shapes.resize(1);
         MaybeTensorProtoToShape(ic, tensor_proto,
                                 &c->output_tensors_as_shapes[0]);
+      } else if (IsCast(node)) {
+        if (c->input_tensors_as_shapes_to_propagate.empty()) {
+          return absl::OkStatus();
+        }
+        const DataType src_type = node.attr().at("SrcT").type();
+        const DataType dst_type = node.attr().at("DstT").type();
+        if ((src_type == DT_INT32 || src_type == DT_INT64) &&
+            (dst_type == DT_INT32 || dst_type == DT_INT64)) {
+          c->output_tensors_as_shapes.resize(1);
+          c->output_tensors_as_shapes[0] =
+              c->input_tensors_as_shapes_to_propagate[0];
+        }
       } else if (IsRank(node)) {
         if (ic->RankKnown(ic->input(0))) {
           // Propagate rank value.
@@ -1731,6 +1860,19 @@ class SymbolicShapeRefiner {
         }
       } else if (IsSize(node)) {
         DimensionHandle size = ic->NumElements(ic->input(0));
+        if (enable_dynamic_value_inference_ && !ic->ValueKnown(size) &&
+            ic->RankKnown(ic->input(0))) {
+          size = ic->MakeDim(1);
+          for (int i = 0; i < ic->Rank(ic->input(0)); ++i) {
+            TF_RETURN_IF_ERROR(
+                ic->Multiply(size, ic->Dim(ic->input(0), i), &size));
+          }
+        }
+        if (enable_dynamic_value_inference_ &&
+            (ic->ValueKnown(size) || ic->GetDimExpr(size) != nullptr)) {
+          c->output_tensors_as_shapes.resize(1);
+          c->output_tensors_as_shapes[0] = ic->MakeShape({size});
+        }
         if (ic->ValueKnown(size)) {
           // Propagate size value.
           int64_t sz = ic->Value(size);
@@ -1751,6 +1893,48 @@ class SymbolicShapeRefiner {
             c->output_tensor_protos[0] = &const_tensors_to_propagate_.back();
           }
         }
+      } else if (enable_dynamic_value_inference_ && op == "Range") {
+        auto scalar_int_value = [&](int input, int64_t* value) {
+          const Tensor* tensor = ic->input_tensor(input);
+          if (tensor == nullptr || tensor->dims() != 0) {
+            return false;
+          }
+          if (tensor->dtype() == DT_INT32) {
+            *value = tensor->scalar<int32>()();
+            return true;
+          }
+          if (tensor->dtype() == DT_INT64) {
+            *value = tensor->scalar<int64_t>()();
+            return true;
+          }
+          return false;
+        };
+
+        int64_t start;
+        int64_t delta;
+        const bool has_positive_constant_delta =
+            scalar_int_value(0, &start) && scalar_int_value(2, &delta) &&
+            start == 0 && delta > 0;
+        if (has_positive_constant_delta &&
+            c->input_tensors_as_shapes_to_propagate.size() > 1) {
+          const ShapeHandle& limit =
+              c->input_tensors_as_shapes_to_propagate[1];
+          if (ic->RankKnown(limit) && ic->Rank(limit) >= 1) {
+            DimensionHandle length = ic->Dim(limit, 0);
+            if (ic->ValueKnown(length) || ic->GetDimExpr(length) != nullptr) {
+              DimensionHandle range_length = length;
+              if (delta > 1) {
+                DimensionHandle adjusted_length;
+                TF_RETURN_IF_ERROR(
+                    ic->Add(length, delta - 1, &adjusted_length));
+                TF_RETURN_IF_ERROR(ic->Divide(
+                    adjusted_length, delta, /*evenly_divisible=*/false,
+                    &range_length));
+              }
+              ic->set_output(0, ic->Vector(range_length));
+            }
+          }
+        }
       } else if (IsShape(node)) {
         c->output_tensors_as_shapes.resize(1);
         c->output_tensors_as_shapes[0] = c->inference_context->input(0);
@@ -1759,7 +1943,126 @@ class SymbolicShapeRefiner {
         for (int i = 0; i < c->inference_context->num_inputs(); ++i) {
           c->output_tensors_as_shapes[i] = c->inference_context->input(i);
         }
-      } else if (node.op() == "ConcatV2") {
+      } else if (IsGather(node)) {
+        if (c->input_tensors_as_shapes_to_propagate.empty()) {
+          return absl::OkStatus();
+        }
+        const ShapeHandle& params = c->input_tensors_as_shapes_to_propagate[0];
+        // Shape-like tensors are propagated here by storing each tensor value
+        // as one dimension of a ShapeHandle. Gather(axis=0) therefore becomes
+        // "pick those propagated entries by index".
+        const Tensor* indices = ic->input_tensor(1);
+        const bool has_axis_input = ic->num_inputs() >= 3;
+        const Tensor* axis = has_axis_input ? ic->input_tensor(2) : nullptr;
+        bool valid = ic->RankKnown(params) && indices != nullptr;
+        int64_t axis_value = 0;
+        if (valid && has_axis_input) {
+          valid = axis != nullptr && axis->NumElements() == 1;
+        }
+        if (valid && has_axis_input) {
+          axis_value = axis->dtype() == DT_INT32 ? axis->scalar<int32>()()
+                                                 : axis->scalar<int64_t>()();
+          valid = (axis_value == 0);
+        }
+        if (valid) {
+          std::vector<DimensionHandle> dims;
+          auto add_index = [&](int64_t raw_index) -> bool {
+            int64_t index = raw_index;
+            if (index < 0) {
+              // Match Gather's negative-index semantics.
+              index += ic->Rank(params);
+            }
+            if (index < 0 || index >= ic->Rank(params)) {
+              return false;
+            }
+            // In this side channel, shape-tensor contents live in the dims.
+            dims.push_back(ic->Dim(params, index));
+            return true;
+          };
+          auto add_indices = [&](const auto& values) -> bool {
+            for (int i = 0; i < values.size(); ++i) {
+              if (!add_index(values(i))) {
+                return false;
+              }
+            }
+            return true;
+          };
+          if (indices->dims() == 0) {
+            valid =
+                add_index(indices->dtype() == DT_INT32 ? indices->scalar<int32>()()
+                                                       : indices->scalar<int64_t>()());
+          } else if (indices->dims() == 1) {
+            if (indices->dtype() == DT_INT32) {
+              valid = add_indices(indices->vec<int32>());
+            } else if (indices->dtype() == DT_INT64) {
+              valid = add_indices(indices->vec<int64_t>());
+            } else {
+              valid = false;
+            }
+          } else {
+            valid = false;
+          }
+          if (valid) {
+            c->output_tensors_as_shapes.resize(1);
+            c->output_tensors_as_shapes[0] = ic->MakeShape(dims);
+          }
+        }
+      } else if (IsProd(node)) {
+        if (c->input_tensors_as_shapes_to_propagate.empty()) {
+          return absl::OkStatus();
+        }
+        const ShapeHandle& input = c->input_tensors_as_shapes_to_propagate[0];
+        // Here `input` represents a propagated shape-vector value, so its
+        // entries live in the dimensions of the ShapeHandle rather than in a
+        // normal tensor buffer.
+        const Tensor* reduction_indices = ic->input_tensor(1);
+        bool keep_dims = false;
+        TF_RETURN_IF_ERROR(GetNodeAttr(node, "keep_dims", &keep_dims));
+        bool valid =
+            ic->RankKnown(input) && reduction_indices != nullptr && !keep_dims;
+        if (valid) {
+          std::vector<int64_t> axes;
+          auto read_axes = [&](const auto& values) {
+            axes.reserve(values.size());
+            for (int i = 0; i < values.size(); ++i) {
+              axes.push_back(values(i));
+            }
+          };
+          if (reduction_indices->dtype() == DT_INT32) {
+            if (reduction_indices->dims() == 0) {
+              axes.push_back(reduction_indices->scalar<int32>()());
+            } else if (reduction_indices->dims() == 1) {
+              read_axes(reduction_indices->vec<int32>());
+            } else {
+              valid = false;
+            }
+          } else if (reduction_indices->dtype() == DT_INT64) {
+            if (reduction_indices->dims() == 0) {
+              axes.push_back(reduction_indices->scalar<int64_t>()());
+            } else if (reduction_indices->dims() == 1) {
+              read_axes(reduction_indices->vec<int64_t>());
+            } else {
+              valid = false;
+            }
+          } else {
+            valid = false;
+          }
+          if (valid) {
+            // For shape-vector contents, Prod(axis=0) means multiply the
+            // propagated entries together.
+            valid = axes.size() == 1 && axes[0] == 0 && ic->Rank(input) > 0;
+          }
+          if (valid) {
+            DimensionHandle product = ic->Dim(input, 0);
+            for (int i = 1; i < ic->Rank(input); ++i) {
+              TF_RETURN_IF_ERROR(ic->Multiply(product, ic->Dim(input, i),
+                                              &product));
+            }
+            c->output_tensors_as_shapes.resize(1);
+            c->output_tensors_as_shapes[0] = ic->MakeShape({product});
+          }
+        }
+      } else if (IsConcat(node)) {
         bool valid = true;
         ShapeHandle result;
         for (int i = 0; i < ic->num_inputs() - 1; ++i) {
@@ -1798,8 +2101,12 @@ class SymbolicShapeRefiner {
             // possible.
             const ShapeHandle& shape_handle =
                 c->input_tensors_as_shapes_to_propagate[i];
-            if (ic->RankKnown(shape_handle) && ic->Rank(shape_handle) >= 1 &&
-                ic->ValueKnown(ic->Dim(shape_handle, 0))) {
+            const bool has_value =
+                ic->RankKnown(shape_handle) && ic->Rank(shape_handle) >= 1 &&
+                (ic->ValueKnown(ic->Dim(shape_handle, 0)) ||
+                 (enable_dynamic_value_inference_ &&
+                  ic->GetDimExpr(ic->Dim(shape_handle, 0)) != nullptr));
+            if (has_value) {
               dims.push_back(ic->Dim(shape_handle, 0));
             } else {
               // This is not from Const, but as it shouldn'be used as symbolic
@@ -1812,6 +2119,26 @@ class SymbolicShapeRefiner {
           c->output_tensors_as_shapes.resize(1);
           c->output_tensors_as_shapes[0] = ic->MakeShape(dims);
         }
+      } else if (IsUnpack(node)) {
+        if (c->input_tensors_as_shapes_to_propagate.empty()) {
+          return absl::OkStatus();
+        }
+        const ShapeHandle& input = c->input_tensors_as_shapes_to_propagate[0];
+        bool valid = ic->RankKnown(input);
+        int axis = 0;
+        if (valid) {
+          TF_RETURN_IF_ERROR(GetNodeAttr(node, "axis", &axis));
+          if (axis < 0) {
+            axis += ic->Rank(input);
+          }
+          valid = (axis == 0 && ic->num_outputs() == ic->Rank(input));
+        }
+        if (valid) {
+          c->output_tensors_as_shapes.resize(ic->num_outputs());
+          for (int i = 0; i < ic->num_outputs(); ++i) {
+            c->output_tensors_as_shapes[i] = ic->MakeShape({ic->Dim(input, i)});
+          }
+        }
       } else if (IsIdentity(node) || IsIdentityNSingleInput(node)) {
         c->output_tensors_as_shapes.resize(1);
         c->output_tensors_as_shapes[0] =
@@ -1919,6 +2246,99 @@ class SymbolicShapeRefiner {
     return absl::OkStatus();
   }
 
+  absl::Status CanonicalizeOutputDims(const NodeDef* node) {
+    NodeContext* ctx = GetNodeContext(node);
+    if (!ctx) return absl::OkStatus();
+
+    InferenceContext* ic = ctx->inference_context.get();
+    for (int out = 0; out < ic->num_outputs(); ++out) {
+      ShapeHandle s = ic->output(out);
+
+      if (!ic->RankKnown(s)) {
+        bool recovered_rank = false;
+        auto it = node->attr().find("_output_shapes");
+        if(it == node->attr().end()){
+          it = node->attr().find("shape");
+        }
+        if (it != node->attr().end() && out < it->second.list().shape_size()) {
+          const TensorShapeProto& proto = it->second.list().shape(out);
+          if (!proto.unknown_rank()) {
+            std::vector<DimensionHandle> dims;
+            dims.reserve(proto.dim_size());
+
+            for (int d = 0; d < proto.dim_size(); ++d) {
+              int64_t size = proto.dim(d).size();
+              if (size >= 0) {
+                dims.push_back(ic->MakeDim(size));
+              } else {
+                dims.push_back(GetUnknownOutputDim(node, out, d));
+              }
+            }
+            s = ic->MakeShape(dims);
+            ic->set_output(out, s);
+            recovered_rank = true;
+          }
+        }
+        if (!recovered_rank && node->op() == "_Arg") {
+          DimensionHandle d0 = GetUnknownOutputDim(node, out, /*dim_id=*/0);
+          ShapeHandle vec = ic->MakeShape({d0});
+          ic->set_output(out, vec);
+          s = vec;
+          recovered_rank = true;
+        }
+
+        if (!recovered_rank) {
+          VLOG(1) << "RANK still unknown. " << node->name();
+          continue;
+        }
+      }
+
+      if (!ic->RankKnown(s)) {
+        continue;
+      }
+
+      bool changed = false;
+      std::vector<DimensionHandle> dims;
+      dims.reserve(ic->Rank(s));
+      for (int d = 0; d < ic->Rank(s); ++d) {
+        DimensionHandle dim = ic->Dim(s, d);
+        const int64_t v = ic->Value(dim);
+        // Keep concrete dims.
+        if (v >= 0) {
+          dims.push_back(dim);
+          continue;
+        }
+        // If already tagged with expr, keep it.
+        auto* dim_expr = ic->GetDimExpr(dim);
+        if (dim_expr != nullptr) {
+          dims.push_back(dim);
+          continue;
+        }
+        auto output_shapes_it = node->attr().find("_output_shapes");
+        if (output_shapes_it != node->attr().end() &&
+            out < output_shapes_it->second.list().shape_size() &&
+            d < output_shapes_it->second.list().shape(out).dim_size()) {
+          const int64_t annotated_size =
+              output_shapes_it->second.list().shape(out).dim(d).size();
+          if (annotated_size >= 0) {
+            changed = true;
+            dims.push_back(ic->MakeDim(annotated_size));
+            continue;
+          }
+        }
+        // Canonicalize ALL unknown dims.
+        DimensionHandle canon = GetUnknownOutputDim(node, out, d);
+        changed |= !dim.SameHandle(canon);
+        dims.push_back(canon);
+      }
+      if (changed) {
+        ShapeHandle new_s = ic->MakeShape(dims);
+        ic->set_output(out, new_s);
+      }
+    }
+    return absl::OkStatus();
+  }
+
   absl::Status InferShapes(const NodeDef& node, NodeContext* c) {
     // Infer the shapes of output tensors.
     if (!c->op_data || c->op_data->shape_inference_fn == nullptr ||
@@ -1940,8 +2360,8 @@ class SymbolicShapeRefiner {
         status.Update(SetUnknownShape(&node, output_port));
       }
     }
-
     // Update NodeContext output fields after shape inference function runs.
+    status.Update(CanonicalizeOutputDims(&node));
     status.Update(MaybeUpdateNodeContextOutput(node, is_fed, c));
 
     return status;
@@ -2048,12 +2468,17 @@ class SymbolicShapeRefiner {
   absl::flat_hash_map<const NodeDef*, NodeContext> node_to_context_;
   absl::flat_hash_map<ShapeId, ShapeHandle> unknown_shapes_;
   absl::flat_hash_map<DimId, DimensionHandle> unknown_dims_;
+  // Stable variable IDs for canonical dimension symbols.
+  absl::flat_hash_map<DimId, int> stable_var_ids_;
+  int next_var_id_ = 1;
   // Store function instantiations only for valid function. If function
   // instantiation failed it will have an `absl::nullopt`.
   absl::flat_hash_map<string, absl::optional<GrapplerFunctionItem>>
       fun_to_grappler_function_item_;
   FunctionLibraryDefinition function_library_;
   const absl::flat_hash_map<string, absl::flat_hash_set<int>>& fed_ports_;
+  const absl::flat_hash_set<string> top_level_input_placeholders_;
+  const int input_arg_shared_symbol_id_;
   // Store TensorProtos for tensor value propagation. Note that we use deque,
   // not vector, as we use pointers to the TensorProtos in this container.
   // Vector may resize and copy the objects into a new buffer, then the existing
@@ -2062,6 +2487,7 @@ class SymbolicShapeRefiner {
 
   // For more aggressive shape and value inference.
   bool aggressive_shape_inference_;
+  bool enable_dynamic_value_inference_;
   ResourceMgr resource_mgr_;
 };
 
@@ -2080,8 +2506,9 @@ class SymbolicShapeManager {
     if (InferenceContext::Rank(s1) > 0 && InferenceContext::Rank(s2) > 0) {
       CHECK_EQ(InferenceContext::Rank(s1), InferenceContext::Rank(s2));
       for (int i = 0; i < InferenceContext::Rank(s1); ++i) {
-        TF_RETURN_IF_ERROR(dims_.Merge(InferenceContext::DimKnownRank(s1, i),
-                                       InferenceContext::DimKnownRank(s2, i)));
+        TF_RETURN_IF_ERROR(
+            MergeDimsWithExpr(InferenceContext::DimKnownRank(s1, i),
+                              InferenceContext::DimKnownRank(s2, i)));
       }
     }
     return absl::OkStatus();
@@ -2090,7 +2517,7 @@ class SymbolicShapeManager {
     if (!d1.IsSet() || !d2.IsSet()) {
       return absl::OkStatus();
     }
-    return dims_.Merge(d1, d2);
+    return MergeDimsWithExpr(d1, d2);
   }
 
   void AsTensorProperties(const ShapeHandle& shape, const DataType& type,
@@ -2104,7 +2531,26 @@ class SymbolicShapeManager {
         shape_inference::DimensionHandle dim =
             InferenceContext::DimKnownRank(actual_shape, j);
         int64_t d = dims_.GetMergedValue(dim);
-        properties->mutable_shape()->add_dim()->set_size(d);
+        TensorShapeProto* output_shape = properties->mutable_shape();
+        auto* out_dim = output_shape->add_dim();
+        out_dim->set_size(d < 0 ? -1 : d);
+        void* root = dims_.RootId(dim);
+        DimExpr* expr = nullptr;
+        if (auto it = dim_root_expr_.find(root); it != dim_root_expr_.end()) {
+          expr = it->second;
+        } else {
+          expr = ExprForDim(dim);
+        }
+        ExpressionProto* output_expr = output_shape->add_expressions();
+        if (expr != nullptr) {
+          DimExprToProto(*expr, output_expr);
+          // TODO: Apply simplification?
+        } else if (d >= 0) {
+          output_expr->set_constant_value(d);
+        } else {
+          DimExprToProto(DimExpr::Unknown(xla::kMissingExpressionSentinel),
+                         output_expr);
+        }
       }
     }
   }
@@ -2132,7 +2578,142 @@ class SymbolicShapeManager {
   }
 
  private:
+  // Get the variable ID from an expression, or -1 if not a variable.
+  static int32_t GetVarId(const DimExpr* e) {
+    if (!e || e->kind() != DimExpr::Kind::kVariable) return -1;
+    return static_cast<const xla::Variable*>(e->get())->get_id();
+  }
+
+  static bool IsConst(const DimExpr* e) {
+    return e && e->kind() == DimExpr::Kind::kConstant;
+  }
+
+  static bool IsVar(const DimExpr* e) {
+    return e && e->kind() == DimExpr::Kind::kVariable;
+  }
+
+  static bool IsPlaceHolder(const DimExpr* e) {
+    if (!e) return false;
+    if (e->kind() != DimExpr::Kind::kVariable) return false;
+    return static_cast<const xla::Variable*>(e->get())->get_id() < 0;
+  }
+
+  static bool IsCompound(const DimExpr* e) {
+    if (!e) return false;
+    switch (e->kind()) {
+      case DimExpr::Kind::kAdd:
+      case DimExpr::Kind::kSub:
+      case DimExpr::Kind::kMul:
+      case DimExpr::Kind::kDiv:
+      case DimExpr::Kind::kMax:
+      case DimExpr::Kind::kGt:
+      case DimExpr::Kind::kSelect:
+        return true;
+      default:
+        return false;
+    }
+  }
+
+  // Ranking: Const > Arg_ > Compound > Var > null
+  static int InfoScore(const DimExpr* e) {
+    if (!e) return 0;
+    if (IsConst(e)) return 4;
+    if (IsPlaceHolder(e)) return 3;
+    if (IsCompound(e)) return 2;
+    if (IsVar(e)) return 1;
+    return 1;  // fallback (shouldn't happen)
+  }
+
+  // Prefer "more informative" but keep deterministic tie-break.
+  static DimExpr* PreferMoreInformative(DimExpr* a, DimExpr* b) {
+    if (a == b) return a;
+    const int sa = InfoScore(a);
+    const int sb = InfoScore(b);
+    if (sa > sb) return a;
+    if (sb > sa) return b;
+    // Same score: keep stable choice.
+    return a;
+  }
+
+  // Get the expr pointer from a dimension handle (accesses private member).
+  static DimExpr* GetExprFromDimHandle(const DimensionHandle& d) {
+    if (!d.IsSet()) return nullptr;
+    return d->expr_;
+  }
+
+  DimExpr* ExprForDim(const DimensionHandle& d) {
+    if (!d.IsSet()) return nullptr;
+    if (DimExpr* expr = GetExprFromDimHandle(d)) {
+      return expr;
+    }
+    if (!InferenceContext::ValueKnown(d)) {
+      return nullptr;
+    }
+    const int64_t value = InferenceContext::Value(d);
+    auto it = const_exprs_.find(value);
+    if (it != const_exprs_.end()) {
+      return it->second.get();
+    }
+    auto expr = std::make_unique<DimExpr>(DimExpr::Const(value));
+    DimExpr* expr_ptr = expr.get();
+    const_exprs_.emplace(value, std::move(expr));
+    return expr_ptr;
+  }
+
+  absl::Status MergeDimsWithExpr(DimensionHandle d1, DimensionHandle d2) {
+    if (!d1.IsSet() || !d2.IsSet()) return absl::OkStatus();
+
+    void* r1 = dims_.RootId(d1);
+    void* r2 = dims_.RootId(d2);
+
+    // Fetch best-known expr for each set.
+    auto get_best = [&](void* r, DimensionHandle d) -> DimExpr* {
+      auto it = dim_root_expr_.find(r);
+      if (it != dim_root_expr_.end()) return it->second;
+      return ExprForDim(d);  // may be null
+    };
+
+    DimExpr* e1 = get_best(r1, d1);
+    DimExpr* e2 = get_best(r2, d2);
+
+    // If already in same UF set, just keep the most informative expr.
+    if (r1 == r2) {
+      DimExpr* existing = nullptr;
+      if (auto it = dim_root_expr_.find(r1); it != dim_root_expr_.end()) {
+        existing = it->second;
+      }
+      DimExpr* chosen = PreferMoreInformative(existing,
+                                            PreferMoreInformative(e1, e2));
+      if (chosen) dim_root_expr_[r1] = chosen;  // keep or upgrade
+      return absl::OkStatus();
+    }
+
+    // Perform UF merge (rank + path compression inside DisjointSet).
+    TF_RETURN_IF_ERROR(dims_.Merge(d1, d2));
+
+    // New root after merge.
+    void* new_root = dims_.RootId(d1);
+
+    // Choose best expr across both sets.
+    DimExpr* chosen = PreferMoreInformative(e1, e2);
+
+    // Remove stale root keys (only the old roots).
+    dim_root_expr_.erase(r1);
+    dim_root_expr_.erase(r2);
+
+    // Preserve any expr already stored at new_root (rare but safe).
+    if (auto it = dim_root_expr_.find(new_root); it != dim_root_expr_.end()) {
+      chosen = PreferMoreInformative(it->second, chosen);
+    }
+
+    if (chosen) dim_root_expr_[new_root] = chosen;
+
+    return absl::OkStatus();
+  }
   DisjointSet<shape_inference::ShapeHandle> shapes_;
+  absl::flat_hash_map<int64_t, std::unique_ptr<DimExpr>> const_exprs_;
+  // Map from union-find root pointer to the best expression for that set.
+  absl::flat_hash_map<void*, DimExpr*> dim_root_expr_;
   DisjointSet<shape_inference::DimensionHandle> dims_;
 };
 
@@ -2518,7 +3099,8 @@ absl::Status GraphProperties::UpdateEnqueue(
 
 absl::Status GraphProperties::InferStatically(
     bool assume_valid_feeds, bool aggressive_shape_inference,
-    bool include_input_tensor_values, bool include_output_tensor_values) {
+    bool include_input_tensor_values, bool include_output_tensor_values,
+    bool enable_dynamic_value_inference) {
   FunctionLibraryDefinition function_library(OpRegistry::Global(),
                                              item_.graph.library());
   absl::flat_hash_map<string, absl::flat_hash_set<int>> fed_ports;
@@ -2606,8 +3188,16 @@ absl::Status GraphProperties::InferStatically(
 
   // Heap-allocate SymbolicShapeRefiner in order to not consume a large amount
   // of stack space.
+  absl::flat_hash_set<string> top_level_input_placeholders;
+  for (const NodeDef& node : item_.graph.node()) {
+    if (IsPlaceholderOp(node.op()) || node.op() == "_Arg") {
+      top_level_input_placeholders.insert(node.name());
+    }
+  }
   auto refiner = std::make_unique<SymbolicShapeRefiner>(
-      graph_view, fed_ports, aggressive_shape_inference);
+      graph_view, fed_ports, aggressive_shape_inference,
+      enable_dynamic_value_inference,
+      std::move(top_level_input_placeholders));
 
   TopoQueue new_shapes(topo_order);
   // Also seed the propagation of shapes in the fanout of primary inputs.
diff --git a/tensorflow/core/grappler/costs/graph_properties.h b/tensorflow/core/grappler/costs/graph_properties.h
index 1d9575e1e5c805..2d3a51cd2975f0 100644
--- a/tensorflow/core/grappler/costs/graph_properties.h
+++ b/tensorflow/core/grappler/costs/graph_properties.h
@@ -99,7 +99,8 @@ class GraphProperties {
   absl::Status InferStatically(bool assume_valid_feeds,
                                bool aggressive_shape_inference,
                                bool include_input_tensor_values,
-                               bool include_output_tensor_values);
+                               bool include_output_tensor_values,
+                               bool enable_dynamic_value_inference = false);
   absl::Status InferStatically(bool assume_valid_feeds,
                                bool aggressive_shape_inference,
                                bool include_tensor_values) {
diff --git a/tensorflow/core/grappler/costs/graph_properties_test.cc b/tensorflow/core/grappler/costs/graph_properties_test.cc
index 53fe8bef1f6f85..d88f9c0cf36033 100644
--- a/tensorflow/core/grappler/costs/graph_properties_test.cc
+++ b/tensorflow/core/grappler/costs/graph_properties_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/tensor.pb.h"  // NOLINT
 #include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/tensor_shape_expr.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/framework/versions.pb.h"
@@ -42,6 +43,40 @@ namespace tensorflow {
 namespace grappler {
 namespace {
 
+DimExpr ShapeDimExpr(const TensorShapeProto& shape, int dim) {
+  if (dim < shape.expressions_size()) {
+    return DimExprFromProto(shape.expressions(dim));
+  }
+  return DimExpr();
+}
+
+bool DimExprEqual(const DimExpr& lhs_expr, const DimExpr& rhs_expr) {
+  if (!lhs_expr || !rhs_expr) return !lhs_expr && !rhs_expr;
+  return lhs_expr == rhs_expr;
+}
+
+bool ShapeDimExprEqual(const TensorShapeProto& lhs, int lhs_dim,
+                       const TensorShapeProto& rhs, int rhs_dim) {
+  return DimExprEqual(ShapeDimExpr(lhs, lhs_dim),
+                      ShapeDimExpr(rhs, rhs_dim));
+}
+
+void ExpectVariableId(const TensorShapeProto& shape, int dim,
+                      int expected_var_id) {
+  ASSERT_GT(shape.dim_size(), dim);
+  const auto& expr = shape.dim(dim).expr();
+  EXPECT_EQ(expr.node_type_case(), ExpressionProto::kVariableId);
+  EXPECT_EQ(expr.variable_id(), expected_var_id);
+}
+
+void ExpectNotVariableId(const TensorShapeProto& shape, int dim,
+                         int unexpected_var_id) {
+  ASSERT_GT(shape.dim_size(), dim);
+  const auto& expr = shape.dim(dim).expr();
+  ASSERT_EQ(expr.node_type_case(), ExpressionProto::kVariableId);
+  EXPECT_NE(expr.variable_id(), unexpected_var_id);
+}
+
 using shape_inference::InferenceContext;
 using shape_inference::ShapeAndType;
 using shape_inference::ShapeHandle;
@@ -776,9 +811,11 @@ TEST_F(GraphPropertiesTest, WhileLoop) {
   // since we concatenated along the batch dim.
   auto shape_in = properties.GetOutputProperties("ones").at(0).shape();
   auto shape_out = properties.GetOutputProperties("while/Exit_1").at(0).shape();
-  EXPECT_GE(-2, shape_in.dim(0).size());
-  EXPECT_GE(-2, shape_out.dim(0).size());
-  EXPECT_NE(shape_in.dim(0).size(), shape_out.dim(0).size());
+  EXPECT_EQ(-1, shape_in.dim(0).size());
+  EXPECT_EQ(-1, shape_out.dim(0).size());
+  EXPECT_TRUE(ShapeDimExpr(shape_in, 0));
+  EXPECT_TRUE(ShapeDimExpr(shape_out, 0));
+  EXPECT_FALSE(ShapeDimExprEqual(shape_in, 0, shape_out, 0));
 }
 
 TEST_F(GraphPropertiesTest, NestedLoop) {
@@ -1969,10 +2006,10 @@ TEST_F(GraphPropertiesTest, SymbolicShapes) {
   const auto shape_c = properties.GetOutputProperties("c").at(0).shape();
   EXPECT_EQ(2, shape_a.dim_size());
   EXPECT_EQ(shape_a.dim_size(), shape_c.dim_size());
-  EXPECT_GE(-2, shape_a.dim(0).size());
-  EXPECT_EQ(shape_a.dim(0).size(), shape_c.dim(0).size());
-  EXPECT_GE(-2, shape_a.dim(1).size());
-  EXPECT_EQ(shape_a.dim(1).size(), shape_c.dim(1).size());
+  EXPECT_EQ(-1, shape_a.dim(0).size());
+  EXPECT_TRUE(ShapeDimExprEqual(shape_a, 0, shape_c, 0));
+  EXPECT_EQ(-1, shape_a.dim(1).size());
+  EXPECT_TRUE(ShapeDimExprEqual(shape_a, 1, shape_c, 1));
 
   PartialTensorShape shape(shape_a);
   EXPECT_FALSE(shape.IsFullyDefined());
@@ -1982,29 +2019,181 @@ TEST_F(GraphPropertiesTest, SymbolicShapes) {
   const auto shape_d = properties.GetOutputProperties("d").at(0).shape();
   EXPECT_EQ(1, shape_b.dim_size());
   EXPECT_EQ(shape_b.dim_size(), shape_d.dim_size());
-  EXPECT_GE(-2, shape_b.dim(0).size());
-  EXPECT_NE(shape_a.dim(0).size(), shape_b.dim(0).size());
-  EXPECT_EQ(shape_b.dim(0).size(), shape_d.dim(0).size());
+  EXPECT_EQ(-1, shape_b.dim(0).size());
+  EXPECT_FALSE(ShapeDimExprEqual(shape_a, 0, shape_b, 0));
+  EXPECT_TRUE(ShapeDimExprEqual(shape_b, 0, shape_d, 0));
 
   const auto shape_e = properties.GetOutputProperties("e").at(0).shape();
   ASSERT_EQ(2, shape_e.dim_size());
-  EXPECT_EQ(shape_e.dim(0).size(), shape_c.dim(0).size());
-  EXPECT_NE(shape_e.dim(1).size(), shape_c.dim(1).size());
-  EXPECT_NE(shape_e.dim(0).size(), shape_d.dim(0).size());
+  EXPECT_TRUE(ShapeDimExprEqual(shape_e, 0, shape_c, 0));
+  EXPECT_TRUE(ShapeDimExprEqual(shape_e, 1, shape_c, 1));
+  EXPECT_FALSE(ShapeDimExprEqual(shape_e, 0, shape_d, 0));
 
   const auto shape_f = properties.GetOutputProperties("f").at(0).shape();
   ASSERT_EQ(2, shape_f.dim_size());
-  EXPECT_EQ(shape_f.dim(0).size(), shape_a.dim(0).size());
-  EXPECT_EQ(shape_f.dim(1).size(), shape_a.dim(1).size());
+  EXPECT_TRUE(ShapeDimExprEqual(shape_f, 0, shape_a, 0));
+  EXPECT_TRUE(ShapeDimExprEqual(shape_f, 1, shape_a, 1));
 
   const auto shape_h = properties.GetOutputProperties("h").at(0).shape();
   ASSERT_EQ(2, shape_f.dim_size());
-  EXPECT_EQ(shape_h.dim(0).size(), shape_c.dim(0).size());
-  EXPECT_EQ(shape_h.dim(1).size(), shape_c.dim(1).size());
+  EXPECT_TRUE(ShapeDimExprEqual(shape_h, 0, shape_c, 0));
+  EXPECT_TRUE(ShapeDimExprEqual(shape_h, 1, shape_c, 1));
 
   const auto shape_j = properties.GetOutputProperties("j").at(0).shape();
   ASSERT_EQ(1, shape_j.dim_size());
-  EXPECT_EQ(shape_j.dim(0).size(), shape_a.dim(1).size());
+  EXPECT_TRUE(ShapeDimExprEqual(shape_j, 0, shape_a, 1));
+}
+
+TEST_F(GraphPropertiesTest, SymbolicInputShapesShareSameSymbol) {
+  GrapplerItem item;
+
+  AttrValue shape_a_attr;
+  TensorShapeProto* shape_a_proto = shape_a_attr.mutable_list()->add_shape();
+  shape_a_proto->add_dim()->set_size(-1);
+  shape_a_proto->add_dim()->set_size(-1);
+  TF_ASSERT_OK(NodeDefBuilder("a", "_Arg")
+                   .Attr("T", DT_FLOAT)
+                   .Attr("index", 0)
+                   .Attr("_output_shapes", shape_a_attr)
+                   .Finalize(item.graph.add_node()));
+
+  AttrValue shape_b_attr;
+  TensorShapeProto* shape_b_proto = shape_b_attr.mutable_list()->add_shape();
+  shape_b_proto->add_dim()->set_size(-1);
+  TF_ASSERT_OK(NodeDefBuilder("b", "_Arg")
+                   .Attr("T", DT_FLOAT)
+                   .Attr("index", 1)
+                   .Attr("_output_shapes", shape_b_attr)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("c", "Identity")
+                   .Input(NodeDefBuilder::NodeOut("a", 0, DT_FLOAT))
+                   .Attr("T", DT_FLOAT)
+                   .Finalize(item.graph.add_node()));
+  TF_ASSERT_OK(NodeDefBuilder("d", "Identity")
+                   .Input(NodeDefBuilder::NodeOut("b", 0, DT_FLOAT))
+                   .Attr("T", DT_FLOAT)
+                   .Finalize(item.graph.add_node()));
+
+  GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(false));
+
+  const auto shape_a = properties.GetOutputProperties("a").at(0).shape();
+  const auto shape_b = properties.GetOutputProperties("b").at(0).shape();
+  const auto shape_c = properties.GetOutputProperties("c").at(0).shape();
+  const auto shape_d = properties.GetOutputProperties("d").at(0).shape();
+
+  ASSERT_EQ(2, shape_a.dim_size());
+  ASSERT_EQ(1, shape_b.dim_size());
+  ExpectVariableId(shape_a, 0, -1);
+  ExpectVariableId(shape_a, 1, -1);
+  ExpectVariableId(shape_b, 0, -1);
+  ExpectVariableId(shape_c, 0, -1);
+  ExpectVariableId(shape_c, 1, -1);
+  ExpectVariableId(shape_d, 0, -1);
+}
+
+TEST_F(GraphPropertiesTest, SymbolicNonInputShapesDoNotShareInputSymbol) {
+  GrapplerItem item;
+
+  AttrValue input_shape_attr;
+  TensorShapeProto* input_shape_proto =
+      input_shape_attr.mutable_list()->add_shape();
+  input_shape_proto->add_dim()->set_size(-1);
+  input_shape_proto->add_dim()->set_size(-1);
+  TF_ASSERT_OK(NodeDefBuilder("input", "_Arg")
+                   .Attr("T", DT_FLOAT)
+                   .Attr("index", 0)
+                   .Attr("_output_shapes", input_shape_attr)
+                   .Finalize(item.graph.add_node()));
+
+  AttrValue reshape_shape_attr;
+  TensorShapeProto* reshape_shape_proto =
+      reshape_shape_attr.mutable_list()->add_shape();
+  reshape_shape_proto->add_dim()->set_size(2);
+  TF_ASSERT_OK(NodeDefBuilder("target_shape", "_Arg")
+                   .Attr("T", DT_INT32)
+                   .Attr("index", 1)
+                   .Attr("_output_shapes", reshape_shape_attr)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("reshaped", "Reshape")
+                   .Input(NodeDefBuilder::NodeOut("input", 0, DT_FLOAT))
+                   .Input(
+                       NodeDefBuilder::NodeOut("target_shape", 0, DT_INT32))
+                   .Attr("T", DT_FLOAT)
+                   .Attr("Tshape", DT_INT32)
+                   .Finalize(item.graph.add_node()));
+
+  GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(false));
+
+  const auto input_shape = properties.GetOutputProperties("input").at(0).shape();
+  const auto reshaped_shape =
+      properties.GetOutputProperties("reshaped").at(0).shape();
+
+  ASSERT_EQ(2, input_shape.dim_size());
+  ASSERT_EQ(2, reshaped_shape.dim_size());
+  ExpectVariableId(input_shape, 0, -1);
+  ExpectVariableId(input_shape, 1, -1);
+  ExpectNotVariableId(reshaped_shape, 0, -1);
+  ExpectNotVariableId(reshaped_shape, 1, -1);
+}
+
+TEST_F(GraphPropertiesTest, SymbolicInputShapesShareOnlyDynamicDims) {
+  GrapplerItem item;
+
+  AttrValue shape_a_attr;
+  TensorShapeProto* shape_a_proto = shape_a_attr.mutable_list()->add_shape();
+  shape_a_proto->add_dim()->set_size(4);
+  shape_a_proto->add_dim()->set_size(-1);
+  TF_ASSERT_OK(NodeDefBuilder("a", "_Arg")
+                   .Attr("T", DT_FLOAT)
+                   .Attr("index", 0)
+                   .Attr("_output_shapes", shape_a_attr)
+                   .Finalize(item.graph.add_node()));
+
+  AttrValue shape_b_attr;
+  TensorShapeProto* shape_b_proto = shape_b_attr.mutable_list()->add_shape();
+  shape_b_proto->add_dim()->set_size(-1);
+  shape_b_proto->add_dim()->set_size(8);
+  TF_ASSERT_OK(NodeDefBuilder("b", "_Arg")
+                   .Attr("T", DT_FLOAT)
+                   .Attr("index", 1)
+                   .Attr("_output_shapes", shape_b_attr)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("c", "Identity")
+                   .Input(NodeDefBuilder::NodeOut("a", 0, DT_FLOAT))
+                   .Attr("T", DT_FLOAT)
+                   .Finalize(item.graph.add_node()));
+  TF_ASSERT_OK(NodeDefBuilder("d", "Identity")
+                   .Input(NodeDefBuilder::NodeOut("b", 0, DT_FLOAT))
+                   .Attr("T", DT_FLOAT)
+                   .Finalize(item.graph.add_node()));
+
+  GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(false));
+
+  const auto shape_a = properties.GetOutputProperties("a").at(0).shape();
+  const auto shape_b = properties.GetOutputProperties("b").at(0).shape();
+  const auto shape_c = properties.GetOutputProperties("c").at(0).shape();
+  const auto shape_d = properties.GetOutputProperties("d").at(0).shape();
+
+  ASSERT_EQ(2, shape_a.dim_size());
+  ASSERT_EQ(2, shape_b.dim_size());
+  ASSERT_EQ(2, shape_c.dim_size());
+  ASSERT_EQ(2, shape_d.dim_size());
+
+  EXPECT_EQ(shape_a.dim(0).size(), 4);
+  EXPECT_EQ(shape_b.dim(1).size(), 8);
+  EXPECT_EQ(shape_c.dim(0).size(), 4);
+  EXPECT_EQ(shape_d.dim(1).size(), 8);
+
+  ExpectVariableId(shape_a, 1, -1);
+  ExpectVariableId(shape_b, 0, -1);
+  ExpectVariableId(shape_c, 1, -1);
+  ExpectVariableId(shape_d, 0, -1);
 }
 
 TEST_F(GraphPropertiesTest, DoNotValidateColocationConstraints) {
@@ -2167,6 +2356,74 @@ TEST_F(GraphPropertiesTest, StridedSlicesOfShapes) {
   EXPECT_EQ(shape_a.dim(1).size(), shape_o2.dim(0).size());
 }
 
+TEST_F(GraphPropertiesTest, SizeContentsPropagateToFillOutput) {
+  tensorflow::Scope scope = tensorflow::Scope::NewRootScope();
+  Output input = ops::Placeholder(
+      scope.WithOpName("input"), DT_FLOAT,
+      ops::Placeholder::Shape(PartialTensorShape({-1, 24})));
+  Output input_size = ops::Size(scope.WithOpName("input_size"), input);
+  Output fill_shape =
+      ops::Stack(scope.WithOpName("fill_shape"), {input_size});
+  Output zero = ops::Const(scope.WithOpName("zero"), 0.0f, {});
+  Output filled = ops::Fill(scope.WithOpName("filled"), fill_shape, zero);
+
+  GrapplerItem item;
+  TF_ASSERT_OK(scope.ToGraphDef(&item.graph));
+
+  GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(
+      /*assume_valid_feeds=*/false,
+      /*aggressive_shape_inference=*/false,
+      /*include_input_tensor_values=*/false,
+      /*include_output_tensor_values=*/false,
+      /*enable_dynamic_value_inference=*/true));
+
+  const TensorShapeProto& inferred_input_shape =
+      properties.GetOutputProperties("input").at(0).shape();
+  const TensorShapeProto& inferred_fill_shape =
+      properties.GetOutputProperties("filled").at(0).shape();
+
+  ASSERT_EQ(1, inferred_fill_shape.dim_size());
+  ASSERT_EQ(2, inferred_input_shape.expressions_size());
+  DimExpr expected =
+      DimExprFromProto(inferred_input_shape.expressions(0)) * 24;
+  EXPECT_TRUE(DimExprEqual(expected, ShapeDimExpr(inferred_fill_shape, 0)));
+}
+
+TEST_F(GraphPropertiesTest, SizeContentsPropagateToRangeOutput) {
+  tensorflow::Scope scope = tensorflow::Scope::NewRootScope();
+  Output input = ops::Placeholder(
+      scope.WithOpName("input"), DT_FLOAT,
+      ops::Placeholder::Shape(PartialTensorShape({-1, 1})));
+  Output input_size = ops::Size(scope.WithOpName("input_size"), input);
+  Output zero = ops::Const(scope.WithOpName("zero"), 0, {});
+  Output three = ops::Const(scope.WithOpName("three"), 3, {});
+  Output range =
+      ops::Range(scope.WithOpName("range"), zero, input_size, three);
+
+  GrapplerItem item;
+  TF_ASSERT_OK(scope.ToGraphDef(&item.graph));
+
+  GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(
+      /*assume_valid_feeds=*/false,
+      /*aggressive_shape_inference=*/false,
+      /*include_input_tensor_values=*/false,
+      /*include_output_tensor_values=*/false,
+      /*enable_dynamic_value_inference=*/true));
+
+  const TensorShapeProto& inferred_input_shape =
+      properties.GetOutputProperties("input").at(0).shape();
+  const TensorShapeProto& inferred_range_shape =
+      properties.GetOutputProperties("range").at(0).shape();
+
+  ASSERT_EQ(1, inferred_range_shape.dim_size());
+  ASSERT_EQ(2, inferred_input_shape.expressions_size());
+  DimExpr expected =
+      (DimExprFromProto(inferred_input_shape.expressions(0)) + 2) / 3;
+  EXPECT_TRUE(DimExprEqual(expected, ShapeDimExpr(inferred_range_shape, 0)));
+}
+
 TEST_F(GraphPropertiesTest, StridedSliceOfShapeWithShrinkAxisMask) {
   tensorflow::Scope scope = tensorflow::Scope::NewRootScope();
   Output placeholder =
@@ -2211,6 +2468,143 @@ TEST_F(GraphPropertiesTest, StridedSliceOfShapeWithShrinkAxisMask) {
   }
 }
 
+TEST_F(GraphPropertiesTest, ShapeTensorContentsThroughGatherProdAndUnpack) {
+  GrapplerItem item;
+
+  TF_ASSERT_OK(NodeDefBuilder("input", "Placeholder")
+                   .Attr("dtype", DT_FLOAT)
+                   .Attr("shape", PartialTensorShape({-1, 26, 8}))
+                   .Finalize(item.graph.add_node()));
+  TF_ASSERT_OK(NodeDefBuilder("shape", "Shape")
+                   .Input("input", 0, DT_FLOAT)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("out_type", DT_INT32)
+                   .Finalize(item.graph.add_node()));
+
+  Tensor gather_indices(DT_INT32, TensorShape({2}));
+  gather_indices.vec<int32>()(0) = 0;
+  gather_indices.vec<int32>()(1) = 1;
+  TF_ASSERT_OK(NodeDefBuilder("gather_indices", "Const")
+                   .Attr("dtype", DT_INT32)
+                   .Attr("value", gather_indices)
+                   .Finalize(item.graph.add_node()));
+
+  Tensor zero_scalar(DT_INT32, TensorShape({}));
+  zero_scalar.scalar<int32>()() = 0;
+  TF_ASSERT_OK(NodeDefBuilder("zero_axis", "Const")
+                   .Attr("dtype", DT_INT32)
+                   .Attr("value", zero_scalar)
+                   .Finalize(item.graph.add_node()));
+
+  Tensor eight_scalar(DT_INT32, TensorShape({}));
+  eight_scalar.scalar<int32>()() = 8;
+  TF_ASSERT_OK(NodeDefBuilder("eight", "Const")
+                   .Attr("dtype", DT_INT32)
+                   .Attr("value", eight_scalar)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("gather", "GatherV2")
+                   .Input("shape", 0, DT_INT32)
+                   .Input("gather_indices", 0, DT_INT32)
+                   .Input("zero_axis", 0, DT_INT32)
+                   .Attr("Tparams", DT_INT32)
+                   .Attr("Tindices", DT_INT32)
+                   .Attr("Taxis", DT_INT32)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("flat_dim", "Prod")
+                   .Input("gather", 0, DT_INT32)
+                   .Input("zero_axis", 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tidx", DT_INT32)
+                   .Attr("keep_dims", false)
+                   .Finalize(item.graph.add_node()));
+
+  std::vector<NodeDefBuilder::NodeOut> gather_pack_inputs(2);
+  gather_pack_inputs[0] = NodeDefBuilder::NodeOut{"flat_dim", 0, DT_INT32};
+  gather_pack_inputs[1] = NodeDefBuilder::NodeOut{"eight", 0, DT_INT32};
+  TF_ASSERT_OK(NodeDefBuilder("gather_shape", "Pack")
+                   .Input(gather_pack_inputs)
+                   .Attr("N", 2)
+                   .Attr("T", DT_INT32)
+                   .Attr("axis", 0)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("reshape_from_gather", "Reshape")
+                   .Input("input", 0, DT_FLOAT)
+                   .Input("gather_shape", 0, DT_INT32)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("Tshape", DT_INT32)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("unpack", "Unpack")
+                   .Input("shape", 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("axis", 0)
+                   .Attr("num", 3)
+                   .Finalize(item.graph.add_node()));
+
+  std::vector<NodeDefBuilder::NodeOut> tail_inputs(2);
+  tail_inputs[0] = NodeDefBuilder::NodeOut{"unpack", 1, DT_INT32};
+  tail_inputs[1] = NodeDefBuilder::NodeOut{"unpack", 2, DT_INT32};
+  TF_ASSERT_OK(NodeDefBuilder("tail_shape", "Pack")
+                   .Input(tail_inputs)
+                   .Attr("N", 2)
+                   .Attr("T", DT_INT32)
+                   .Attr("axis", 0)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("tail_prod", "Prod")
+                   .Input("tail_shape", 0, DT_INT32)
+                   .Input("zero_axis", 0, DT_INT32)
+                   .Attr("T", DT_INT32)
+                   .Attr("Tidx", DT_INT32)
+                   .Attr("keep_dims", false)
+                   .Finalize(item.graph.add_node()));
+
+  std::vector<NodeDefBuilder::NodeOut> unpack_pack_inputs(2);
+  unpack_pack_inputs[0] = NodeDefBuilder::NodeOut{"unpack", 0, DT_INT32};
+  unpack_pack_inputs[1] = NodeDefBuilder::NodeOut{"tail_prod", 0, DT_INT32};
+  TF_ASSERT_OK(NodeDefBuilder("unpack_shape", "Pack")
+                   .Input(unpack_pack_inputs)
+                   .Attr("N", 2)
+                   .Attr("T", DT_INT32)
+                   .Attr("axis", 0)
+                   .Finalize(item.graph.add_node()));
+
+  TF_ASSERT_OK(NodeDefBuilder("reshape_from_unpack", "Reshape")
+                   .Input("input", 0, DT_FLOAT)
+                   .Input("unpack_shape", 0, DT_INT32)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("Tshape", DT_INT32)
+                   .Finalize(item.graph.add_node()));
+
+  GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(true));
+
+  const auto& input_shape = properties.GetOutputProperties("input")[0].shape();
+  const auto& gather_reshape_shape =
+      properties.GetOutputProperties("reshape_from_gather")[0].shape();
+  const auto& unpack_reshape_shape =
+      properties.GetOutputProperties("reshape_from_unpack")[0].shape();
+
+  ASSERT_EQ(3, input_shape.dim_size());
+  ASSERT_EQ(2, gather_reshape_shape.dim_size());
+  ASSERT_EQ(2, unpack_reshape_shape.dim_size());
+  ASSERT_TRUE(ShapeDimExpr(input_shape, 0));
+  ASSERT_TRUE(ShapeDimExpr(gather_reshape_shape, 0));
+  ASSERT_TRUE(ShapeDimExpr(unpack_reshape_shape, 0));
+
+  DimExpr input_expr = ShapeDimExpr(input_shape, 0);
+  auto expected_gather_dim0 = input_expr * DimExpr::Const(26);
+  EXPECT_TRUE(DimExprEqual(expected_gather_dim0,
+                           ShapeDimExpr(gather_reshape_shape, 0)));
+  EXPECT_EQ(8, gather_reshape_shape.dim(1).size());
+
+  EXPECT_TRUE(ShapeDimExprEqual(input_shape, 0, unpack_reshape_shape, 0));
+  EXPECT_EQ(208, unpack_reshape_shape.dim(1).size());
+}
+
 TEST_F(GraphPropertiesTest, ValuePropagationThroughArithmeticOps) {
   tensorflow::Scope s = tensorflow::Scope::NewRootScope();
   Output a = ops::Const(s.WithOpName("a"), {5, 7}, {2});
@@ -2255,6 +2649,37 @@ TEST_F(GraphPropertiesTest, ValuePropagationThroughArithmeticOps) {
   ExpectTensorValues({20, 24}, c_plus_b_plus_2a_prop.value());
 }
 
+TEST_F(GraphPropertiesTest,
+       DynamicArithmeticFallsBackToTensorValueInference) {
+  tensorflow::Scope scope = tensorflow::Scope::NewRootScope();
+  Output one = ops::Const(scope.WithOpName("one"), {1}, {1});
+  Output two = ops::Const(scope.WithOpName("two"), {2}, {1});
+  Output unknown = ops::Const(scope.WithOpName("unknown"), {-1}, {1});
+  Output negative = ops::Sub(scope.WithOpName("negative"), one, two);
+  Output unknown_plus_one =
+      ops::Add(scope.WithOpName("unknown_plus_one"), unknown, one);
+
+  GrapplerItem item;
+  TF_ASSERT_OK(scope.ToGraphDef(&item.graph));
+  GraphProperties properties(item);
+  TF_ASSERT_OK(properties.InferStatically(
+      /*assume_valid_feeds=*/false,
+      /*aggressive_shape_inference=*/true,
+      /*include_input_tensor_values=*/true,
+      /*include_output_tensor_values=*/true,
+      /*enable_dynamic_value_inference=*/true));
+
+  const auto& negative_prop =
+      properties.GetOutputProperties("negative").at(0);
+  ASSERT_TRUE(negative_prop.has_value());
+  ExpectTensorValues({-1}, negative_prop.value());
+
+  const auto& unknown_plus_one_prop =
+      properties.GetOutputProperties("unknown_plus_one").at(0);
+  ASSERT_TRUE(unknown_plus_one_prop.has_value());
+  ExpectTensorValues({0}, unknown_plus_one_prop.value());
+}
+
 TEST_F(GraphPropertiesTest, ShapeAnnotation) {
   GrapplerItem item;
   TF_ASSERT_OK(NodeDefBuilder("Input", "Placeholder")
diff --git a/tensorflow/core/grappler/optimizers/remapper.cc b/tensorflow/core/grappler/optimizers/remapper.cc
index a5a48347f07517..c468801a6a74db 100644
--- a/tensorflow/core/grappler/optimizers/remapper.cc
+++ b/tensorflow/core/grappler/optimizers/remapper.cc
@@ -741,6 +741,16 @@ bool IsBiasSemanticAdd(const RemapperContext& ctx,
   return false;
 }
 
+bool MaybeCopyOutputShapesAttr(const NodeDef& from, NodeDef* fused_op) {
+  const string output_shape_attr_name = "_output_shapes";
+  if (from.attr().count(output_shape_attr_name) > 0) {
+    auto shape_attrs = from.attr().at(output_shape_attr_name);
+    AddNodeAttr(output_shape_attr_name, shape_attrs, fused_op);
+    return true;
+  }
+  return false;
+}
+
 void AddInputShapesAttr(const RemapperContext& ctx, int node_index) {
   auto mutable_node = ctx.graph_view.graph()->mutable_node(node_index);
 
@@ -3334,6 +3344,7 @@ absl::Status AddFusedContractionNode(RemapperContext* ctx,
   }
 
   SetFusedOpAttributes(&fused_op, {"BiasAdd"});
+  MaybeCopyOutputShapesAttr(bias_add, &fused_op);
   utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder();
   absl::Status status;
   mutation->AddNode(std::move(fused_op), &status);
@@ -3439,6 +3450,7 @@ absl::Status AddFusedContractionNode(
   }
 
   SetFusedOpAttributes(&fused_op, {"BiasAdd", activation.op()});
+  MaybeCopyOutputShapesAttr(activation, &fused_op);
 
   utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder();
   absl::Status status;
@@ -3630,6 +3642,7 @@ absl::Status AddFusedContractionNode(
   }
 
   SetFusedOpAttributes(&contraction_node, {"BiasAdd", "Add"}, 2);
+  MaybeCopyOutputShapesAttr(add, &contraction_node);
 
   utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder();
   absl::Status status;
@@ -3729,6 +3742,7 @@ absl::Status AddFusedContractionNode(
   }
 
   SetFusedOpAttributes(&fused_conv, {"BiasAdd", "Add", activation.op()}, 2);
+  MaybeCopyOutputShapesAttr(add, &fused_conv);
 
   utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder();
   absl::Status status;
diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc
index 864855de1d69f6..8a5936980ad578 100644
--- a/tensorflow/core/kernels/function_ops.cc
+++ b/tensorflow/core/kernels/function_ops.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/gradients.h"
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/common_runtime/memory_types.h"
+#include "tensorflow/core/framework/batch_size_resource.h"
 #include "tensorflow/core/framework/cancellation.h"
 #include "tensorflow/core/framework/full_type.pb.h"
 #include "tensorflow/core/framework/full_type_util.h"
@@ -42,6 +43,13 @@ static constexpr const char* const kGradientOp =
 ArgOp::ArgOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
   OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_));
   OP_REQUIRES_OK(ctx, ctx->GetAttr("index", &index_));
+
+  Status s = ctx->GetAttr("_dynamic_dim", &dynamic_dim_);
+  if (IsNotFound(s)) {
+    dynamic_dim_ = -1;
+  } else {
+    OP_REQUIRES_OK(ctx, s);
+  }
 }
 
 void ArgOp::Compute(OpKernelContext* ctx) {
@@ -59,16 +67,48 @@ void ArgOp::Compute(OpKernelContext* ctx) {
     }
   };
 
+  Tensor t;
+  int64_t batch_size = -1;
   if (frame->CanConsumeArg(index_)) {
-    Tensor val;
-    frame->ConsumeArg(index_, &val);
-    OP_REQUIRES_OK(ctx, validate_type(val));
-    ctx->set_output(0, std::move(val));
+    frame->ConsumeArg(index_, &t);
+    OP_REQUIRES_OK(ctx, validate_type(t));
+    if (dynamic_dim_ >= 0) {
+      batch_size = t.dim_size(dynamic_dim_);
+    }
+    ctx->set_output(0, std::move(t));
   } else {
     OP_REQUIRES_OK(ctx, frame->GetArg(index_, &val));
     OP_REQUIRES_OK(ctx, validate_type(*val));
     ctx->set_output(0, *val);
   }
+  if (dynamic_dim_ >= 0) {
+    BatchSizeResource* bsr = nullptr;
+    ScopedStepContainer* step_container = ctx->step_container();
+
+    OP_REQUIRES_OK(ctx, step_container->LookupOrCreate<BatchSizeResource>(
+                            ctx->resource_manager(), BatchSizeResourceName, &bsr,
+                            [](BatchSizeResource** ret) -> Status {
+                              *ret = new BatchSizeResource();
+                              return OkStatus();
+                            }));
+
+    if (batch_size < 0) {
+      batch_size = val->dim_size(dynamic_dim_);
+    }
+    VLOG(1) << "Found batch_size in dimension #" << dynamic_dim_;
+    if (bsr->GetBatchSize() == 0) {
+      bsr->SetBatchSize(batch_size);
+      VLOG(1) << "Set batch_size from 0 to " << batch_size
+              << ". step_id: " << ctx->step_id();
+    } else if (bsr->GetBatchSize() != batch_size) {
+      VLOG(1) << "Warning: Set batch_size from " << bsr->GetBatchSize()
+              << ". step_id: " << ctx->step_id();
+      bsr->SetBatchSize(batch_size);
+    } else {
+      VLOG(1) << "batch_size already set to " << batch_size;
+    }
+    bsr->Unref();
+  }
 }
 
 RetvalOp::RetvalOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
diff --git a/tensorflow/core/kernels/function_ops.h b/tensorflow/core/kernels/function_ops.h
index 552e1e6c515e3b..fed05ead2007ad 100644
--- a/tensorflow/core/kernels/function_ops.h
+++ b/tensorflow/core/kernels/function_ops.h
@@ -38,6 +38,7 @@ class ArgOp : public OpKernel {
  private:
   int index_;
   DataType dtype_;
+  int dynamic_dim_;
 
   ArgOp(const ArgOp&) = delete;
   void operator=(const ArgOp&) = delete;
@@ -54,6 +55,7 @@ class RetvalOp : public OpKernel {
  private:
   int index_;
   DataType dtype_;
+  int dynamic_dim_;
 
   RetvalOp(const RetvalOp&) = delete;
   void operator=(const RetvalOp&) = delete;
diff --git a/tensorflow/core/kernels/padding_fifo_queue.cc b/tensorflow/core/kernels/padding_fifo_queue.cc
index 3b50099fb9997c..ed669f60550501 100644
--- a/tensorflow/core/kernels/padding_fifo_queue.cc
+++ b/tensorflow/core/kernels/padding_fifo_queue.cc
@@ -400,7 +400,14 @@ std::vector<TensorShape> PaddingFIFOQueue::ConvertShapesPartialDimensionsToZero(
   for (size_t i = 0; i < shapes.size(); ++i) {
     const PartialTensorShape& partial = partial_shapes[i];
     TensorShape& shape = shapes[i];
-    for (int64_t s : partial.dim_sizes()) shape.AddDim(s < 0 ? 0 : s);
+    for (int d = 0; d < partial.dims(); ++d) {
+      shape.AddDim(partial.dim_size(d) < 0 ? 0 : partial.dim_size(d));
+      xla::DExpr expr = partial.get_filled_expression(d);
+      if (expr && expr->is_constant() && expr->get_val() < 0) {
+        expr = xla::DExpr::Const(0);
+      }
+      shape.AddExpression(expr);
+    }
   }
   return shapes;
 }
diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc
index 4523126f9a71bd..7a2cf9d55922af 100644
--- a/tensorflow/core/kernels/strided_slice_op.cc
+++ b/tensorflow/core/kernels/strided_slice_op.cc
@@ -309,6 +309,8 @@ class StridedSliceAssignOp : public OpKernel {
     bool is_simple_slice = true;
     absl::InlinedVector<int64_t, 4UL> begin;
     absl::InlinedVector<int64_t, 4UL> end;
+    absl::InlinedVector<xla::DExpr, 4UL> begin_expr;
+    absl::InlinedVector<xla::DExpr, 4UL> end_expr;
     absl::InlinedVector<int64_t, 4UL> strides;
 
     Tensor* old_lhs = nullptr;
@@ -353,7 +355,7 @@ class StridedSliceAssignOp : public OpKernel {
                      old_lhs->shape(), begin_mask, end_mask, ellipsis_mask,
                      new_axis_mask, shrink_axis_mask, &processing_shape,
                      &final_shape, &is_identity, &is_simple_slice, &slice_dim0,
-                     &begin, &end, &strides, &shape_spec));
+                     &begin, &end, &strides, &begin_expr, &end_expr, &shape_spec));
 
     if (processing_shape.num_elements() > 0) {
       const Tensor& input = context->input(4);
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 8d53c6dbb38425..b570286261f75b 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -266,15 +266,38 @@ absl::Status SetOutputShapeForReshape(InferenceContext* c) {
                                      true /* evenly_divisible */,
                                      &inferred_dim));
         DimensionHandle unknown_in_dim = c->Dim(in, in_unknown_idx);
+        // Record the inferred scale factor on the input-side unknown.
         TF_RETURN_IF_ERROR(
             c->Merge(unknown_in_dim, inferred_dim, &unknown_in_dim));
       } else if (in_unknown_idx >= 0 && out_unknown_idx >= 0) {
         // Exactly one unknown dimension in both input and output. These 2 are
-        // equal iff the known elements are equal.
+        // related through the ratio of known elements.
+        DimensionHandle unknown_in_dim = c->Dim(in, in_unknown_idx);
         if (c->Value(known_in_elems) == c->Value(known_out_elems)) {
-          DimensionHandle unknown_in_dim = c->Dim(in, in_unknown_idx);
+          // Example: [?,8] -> [?,8], so the unknown is unchanged.
           TF_RETURN_IF_ERROR(
               c->ReplaceDim(out, out_unknown_idx, unknown_in_dim, &out));
+        } else if (c->Value(known_out_elems) > 0 &&
+                   c->Value(known_in_elems) % c->Value(known_out_elems) == 0) {
+          DimensionHandle inferred_out_dim;
+          // Example: [?,26,8] -> [?,8], so the output unknown absorbs *26.
+          TF_RETURN_IF_ERROR(c->Multiply(
+              unknown_in_dim,
+              c->Value(known_in_elems) / c->Value(known_out_elems),
+              &inferred_out_dim));
+          TF_RETURN_IF_ERROR(
+              c->ReplaceDim(out, out_unknown_idx, inferred_out_dim, &out));
+        } else if (c->Value(known_in_elems) > 0 &&
+                   c->Value(known_out_elems) % c->Value(known_in_elems) == 0) {
+          DimensionHandle inferred_out_dim;
+          // Example: [?,8] -> [?,26,8], so the output unknown is the
+          // input-side unknown divided by 26.
+          TF_RETURN_IF_ERROR(c->Divide(
+              unknown_in_dim,
+              c->Value(known_out_elems) / c->Value(known_in_elems),
+              true /* evenly_divisible */, &inferred_out_dim));
+          TF_RETURN_IF_ERROR(
+              c->ReplaceDim(out, out_unknown_idx, inferred_out_dim, &out));
         }
       }
     }
@@ -719,6 +742,8 @@ REGISTER_OP("SplitV")
 REGISTER_OP("Const")
     .Output("output: dtype")
     .Attr("value: tensor")
+    .Attr("user_inferred_shape: shape = {}")
+    .Attr("has_dynamic: bool = false")
     .Attr("dtype: type")
     .SetShapeFn([](InferenceContext* c) {
       const TensorProto* proto = nullptr;
diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc
index 93c5a7e9818ae2..a2bb25f9c2e923 100644
--- a/tensorflow/core/util/strided_slice_op.cc
+++ b/tensorflow/core/util/strided_slice_op.cc
@@ -22,6 +22,8 @@ limitations under the License.
 
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "xla/shape.h"
+
 
 namespace tensorflow {
 namespace {
@@ -55,6 +57,8 @@ struct StridedSliceDenseSpec {
   bool end_valid;
   absl::InlinedVector<int64_t, 4UL>& begin;
   absl::InlinedVector<int64_t, 4UL>& end;
+  absl::InlinedVector<xla::DExpr, 4UL>& begin_expr;
+  absl::InlinedVector<xla::DExpr, 4UL>& end_expr;
   absl::InlinedVector<int64_t, 4UL>& strides;
   // This vector helps construct the final shape of the slice.
   // The final tensor is reduced in rank whenever a single index e.g. foo[3]
@@ -97,6 +101,8 @@ static absl::Status BuildDenseSpec(const StridedSliceSparseSpec& sparse,
   // to remove any ellipsis
   dense->begin.resize(dense->dims);
   dense->end.resize(dense->dims);
+  dense->begin_expr.resize(dense->dims);
+  dense->end_expr.resize(dense->dims);
   dense->strides.resize(dense->dims);
   dense->input_shape_gather_indices_sparse.resize(dense->dims);
   // What indices to get the final shape from.
@@ -127,6 +133,8 @@ static absl::Status BuildDenseSpec(const StridedSliceSparseSpec& sparse,
         for (; full_index < next_index; full_index++) {
           // new_axis' aren't real axis so you have to skip
           dense->begin[full_index] = dense->end[full_index] = 0;
+          dense->begin_expr[full_index] = dense->end_expr[full_index] =
+              xla::DExpr::Const(0);
           dense->strides[full_index] = 1;
           dense->begin_mask |= (1 << full_index);
           dense->end_mask |= (1 << full_index);
@@ -150,9 +158,13 @@ static absl::Status BuildDenseSpec(const StridedSliceSparseSpec& sparse,
         // Gather slicing spec into appropriate index
         if (begin_flat != nullptr) {
           dense->begin[full_index] = internal::SubtleMustCopy<T>(begin_flat[i]);
+          dense->begin_expr[full_index] =
+              xla::DExpr::Const(dense->begin[full_index]);
         }
         if (end_flat != nullptr) {
           dense->end[full_index] = internal::SubtleMustCopy<T>(end_flat[i]);
+          dense->end_expr[full_index] =
+              xla::DExpr::Const(dense->end[full_index]);
         }
         dense->strides[full_index] =
             internal::SubtleMustCopy<T>(strides_flat[i]);
@@ -193,10 +205,29 @@ absl::Status ValidateStridedSliceOp(
     absl::InlinedVector<int64_t, 4UL>* begin,
     absl::InlinedVector<int64_t, 4UL>* end,
     absl::InlinedVector<int64_t, 4UL>* strides,
+    absl::InlinedVector<xla::DExpr, 4UL>* begin_expr,
+    absl::InlinedVector<xla::DExpr, 4UL>* end_expr,
     StridedSliceShapeSpec* shape_spec) {
+  absl::InlinedVector<xla::DExpr, 4UL> b;
+  absl::InlinedVector<xla::DExpr, 4UL> e;
+
+  // HACK
+  if (begin_expr == nullptr) {
+    for (int i : *begin) {
+      b.push_back(xla::DExpr::Const(i));
+    }
+    begin_expr = &b;
+  }
+  if (end_expr == nullptr) {
+    for (int i : *end) {
+      e.push_back(xla::DExpr::Const(i));
+    }
+    end_expr = &e;
+  }
+
   if (input_shape.unknown_rank()) {
     // Note: If the rank is unknown, "input_shape.dims()" is -1.
-    return errors::InvalidArgument("Unexpected input_shape with unknown rank");
+  return errors::InvalidArgument("Unexpected input_shape with unknown rank");
   }
 
   const bool begin_is_wrong =
@@ -271,6 +302,7 @@ absl::Status ValidateStridedSliceOp(
   // we need to produce the missing begin_mask for the first two
   // dimensions i.e. from begin_mask_spec=0, end_mask_spec=2
   // we achieve begin_mask=6, end_mask=7
+
   StridedSliceDenseSpec dense_spec = {input_shape.dims(),
                                       0 /* begin_mask */,
                                       0 /* end_mask */,
@@ -278,6 +310,8 @@ absl::Status ValidateStridedSliceOp(
                                       false /* end_valid */,
                                       *begin,
                                       *end,
+                                      *begin_expr,
+                                      *end_expr,
                                       *strides};
 
   if (strides_tensor.dtype() == DT_INT32) {
@@ -296,17 +330,24 @@ absl::Status ValidateStridedSliceOp(
   *slice_dim0 = true;
   *is_simple_slice = true;
   processing_shape->Clear();
+  auto dim_exprs = input_shape.get_filled_expressions();
   for (int i = 0; i < input_shape.dims(); ++i) {
     int64_t& begin_i = (*begin)[i];
     int64_t& end_i = (*end)[i];
     int64_t& stride_i = (*strides)[i];
     int64_t dim_i = input_shape.dim_size(i);
+
+    xla::DExpr dim_i_expr =
+        i < dim_exprs.size() ? dim_exprs[i] : xla::DExpr::Const(dim_i);
+
     if (stride_i == 0) {
       return errors::InvalidArgument("strides[", i, "] must be non-zero");
     }
     bool shrink_i = (dense_spec.shrink_axis_mask & (1 << i));
     if (dim_i == -1) {
       processing_shape->AddDim(shrink_i ? 1 : -1);
+      processing_shape->AddExpression(shrink_i ? xla::DExpr::Const(1)
+                                               : xla::DExpr::Const(-1));
       continue;
     }
 
@@ -315,15 +356,37 @@ absl::Status ValidateStridedSliceOp(
     const std::array<int64_t, 2> valid_range = {
         {stride_i > 0 ? 0 : -1, stride_i > 0 ? dim_i : dim_i - 1}};
 
+    const std::array<xla::DExpr, 2> valid_range_expr = {
+        {stride_i > 0 ? xla::DExpr::Const(0) : xla::DExpr::Const(-1),
+         stride_i > 0 ? dim_i_expr
+                       : (dim_i_expr - xla::DExpr::Const(1)).simplify()}};
+
     auto canonical = [stride_i, dim_i, masks, valid_range](int64_t x, int c) {
       if (masks[c]) {
         return stride_i > 0 ? valid_range[c] : valid_range[(c + 1) & 1];
       } else {
         int64_t x_fwd =
             x < 0 ? dim_i + x : x;  // make negative indices positive
-        return x_fwd < valid_range[0]
-                   ? valid_range[0]
-                   : x_fwd > valid_range[1] ? valid_range[1] : x_fwd;
+        return x_fwd < valid_range[0]   ? valid_range[0]
+               : x_fwd > valid_range[1] ? valid_range[1]
+                                        : x_fwd;
+      }
+    };
+    auto canonical_expr = [stride_i, dim_i, masks, valid_range,
+                           valid_range_expr, dim_i_expr](int64_t x, int c) {
+      if (masks[c]) {
+        return stride_i > 0 ? valid_range_expr[c]
+                            : valid_range_expr[(c + 1) & 1];
+      } else {
+        int64_t x_fwd =
+            x < 0 ? dim_i + x : x;  // make negative indices positive
+        xla::DExpr x_expr = xla::DExpr::Const(x);
+        xla::DExpr x_fwd_expr =
+            x < 0 ? (dim_i_expr + x_expr)
+                  : x_expr;  // make negative indices positive
+        return x_fwd < valid_range[0]   ? valid_range_expr[0]
+               : x_fwd > valid_range[1] ? valid_range_expr[1]
+                                        : x_fwd_expr;
       }
     };
     if (shrink_i && stride_i <= 0) {
@@ -341,15 +404,30 @@ absl::Status ValidateStridedSliceOp(
         // and canonical puts these to n-1 and 0, which implies a degenerate
         // interval. Fortunately, it is now safe to re-create end as begin+1.
         int64_t x_fwd = begin_i < 0 ? dim_i + begin_i : begin_i;
+        xla::DExpr x_fwd_expr = begin_i < 0
+                                    ? (dim_i_expr + (*begin_expr)[i]).simplify()
+                                    : (*begin_expr)[i];
         begin_i = x_fwd;
         end_i = begin_i + 1;
+
+        (*begin_expr)[i] = x_fwd_expr;
+        (*end_expr)[i] = ((*begin_expr)[i] + xla::DExpr::Const(1)).simplify();
+
         if (x_fwd < 0 || x_fwd >= dim_i) {
           return errors::InvalidArgument(
               "slice index ", begin_i, " of dimension ", i, " out of bounds.");
         }
       } else {
-        begin_i = canonical(begin_i, 0);
-        end_i = canonical(end_i, 1);
+        const int64_t begin_raw = begin_i;
+        const int64_t end_raw = end_i;
+        begin_i = canonical(begin_raw, 0);
+        end_i = canonical(end_raw, 1);
+        if (begin_expr) {
+          (*begin_expr)[i] = canonical_expr(begin_raw, 0).simplify();
+        }
+        if (end_expr) {
+          (*end_expr)[i] = canonical_expr(end_raw, 1).simplify();
+        }
       }
       // Update optimization values
       bool take_all_in_dimension =
@@ -362,14 +440,17 @@ absl::Status ValidateStridedSliceOp(
     }
     // Compute the processing shape (the intermediate Eigen will produce)
     int64_t interval_length;
+    xla::DExpr interval_length_expr;
     bool known_interval = false;
     if (dense_spec.begin_valid && dense_spec.end_valid) {
       interval_length = end_i - begin_i;
+      interval_length_expr = ((*end_expr)[i] - (*begin_expr)[i]).simplify();
       known_interval = true;
     } else if (shrink_i) {
       // The dimension is still known as 1 for the processing_shape, but will be
       // discarded for the final shape.
       interval_length = 1;
+      interval_length_expr = xla::DExpr::Const(1);
       known_interval = true;
     } else if (begin_and_end_masked) {
       // Even if we don't have values for begin or end, we do know that this
@@ -378,25 +459,35 @@ absl::Status ValidateStridedSliceOp(
       if (dim_i >= 0) {
         if (stride_i < 0) {
           interval_length = -dim_i;
+          interval_length_expr = (xla::DExpr::Const(-1) * dim_i_expr).simplify();
         } else {
           interval_length = dim_i;
+          interval_length_expr = dim_i_expr;
         }
         known_interval = true;
       }
     }
     if (known_interval) {
       int64_t size_i;
+      xla::DExpr size_i_expr;
       // Hold zero if the interval is degenerate, otherwise account for
       // remainder
       if (interval_length == 0 || ((interval_length < 0) != (stride_i < 0))) {
         size_i = 0;
+        size_i_expr = xla::DExpr::Const(0);
       } else {
         size_i = interval_length / stride_i +
                  (interval_length % stride_i != 0 ? 1 : 0);
+        size_i_expr =
+            (interval_length_expr / xla::DExpr::Const(stride_i)) +
+            (interval_length % stride_i != 0 ? xla::DExpr::Const(1)
+                                             : xla::DExpr::Const(0));
       }
       processing_shape->AddDim(size_i);
+      processing_shape->AddExpression(size_i_expr.simplify());
     } else {
       processing_shape->AddDim(-1);
+      processing_shape->AddExpression(xla::DExpr::Const(-1));
     }
   }
 
@@ -425,12 +516,15 @@ absl::Status ValidateStridedSliceOp(
         dense_spec.final_shape_gather_indices_sparse[dense_dim];
     if (gather_index >= 0) {
       final_shape->AddDim(processing_shape->dim_size(gather_index));
+      final_shape->AddExpression(
+          processing_shape->get_filled_expression(gather_index));
       if (shape_spec != nullptr) {
         shape_spec->output_to_sparse_mapping.push_back(sparse_index);
         shape_spec->output_to_processing_mapping.push_back(gather_index);
       }
     } else if (gather_index == kNewAxis) {
       final_shape->AddDim(1);
+      final_shape->AddExpression(xla::DExpr::Const(1));
       if (shape_spec != nullptr) {
         shape_spec->output_to_sparse_mapping.push_back(-1);
         shape_spec->output_to_processing_mapping.push_back(-1);
@@ -441,6 +535,7 @@ absl::Status ValidateStridedSliceOp(
   return absl::OkStatus();
 }
 
+
 absl::Status ValidateStridedSliceOp(
     const Tensor* begin_tensor, const Tensor* end_tensor,
     const Tensor& strides_tensor, const PartialTensorShape& input_shape,
@@ -451,16 +546,16 @@ absl::Status ValidateStridedSliceOp(
     absl::InlinedVector<int64_t, 4UL>* begin,
     absl::InlinedVector<int64_t, 4UL>* end,
     absl::InlinedVector<int64_t, 4UL>* strides,
+    absl::InlinedVector<xla::DExpr, 4UL>* begin_expr,
+    absl::InlinedVector<xla::DExpr, 4UL>* end_expr,
     StridedSliceShapeSpec* shape_spec) {
-  // Validate with PartialTensorShape output
   PartialTensorShape partial_processing_shape, partial_final_shape;
   TF_RETURN_IF_ERROR(ValidateStridedSliceOp(
       begin_tensor, end_tensor, strides_tensor, input_shape, begin_mask_spec,
       end_mask_spec, ellipsis_mask, new_axis_mask, shrink_axis_mask,
       &partial_processing_shape, &partial_final_shape, is_identity,
-      is_simple_slice, slice_dim0, begin, end, strides, shape_spec));
-
-  // Verify that the output shapes are fully known
+      is_simple_slice, slice_dim0, begin, end, strides, begin_expr, end_expr,
+      shape_spec));
   if (!partial_processing_shape.AsTensorShape(processing_shape) ||
       !partial_final_shape.AsTensorShape(final_shape)) {
     return errors::Internal("ValidateStridedSliceOp returned partial shapes ",
diff --git a/tensorflow/core/util/strided_slice_op.h b/tensorflow/core/util/strided_slice_op.h
index 356b77a2a0f5b6..c7240d15adc6c7 100644
--- a/tensorflow/core/util/strided_slice_op.h
+++ b/tensorflow/core/util/strided_slice_op.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "xla/shape.h"
 
 namespace tensorflow {
 
@@ -74,6 +75,8 @@ absl::Status ValidateStridedSliceOp(
     absl::InlinedVector<int64_t, 4UL>* begin,
     absl::InlinedVector<int64_t, 4UL>* end,
     absl::InlinedVector<int64_t, 4UL>* strides,
+    absl::InlinedVector<xla::DExpr, 4UL>* begin_expr = nullptr,
+    absl::InlinedVector<xla::DExpr, 4UL>* end_expr = nullptr,
     StridedSliceShapeSpec* shape_spec = nullptr);
 
 // Same as above, but the outputs are TensorShape, not PartialTensorShape
@@ -87,6 +90,8 @@ absl::Status ValidateStridedSliceOp(
     absl::InlinedVector<int64_t, 4UL>* begin,
     absl::InlinedVector<int64_t, 4UL>* end,
     absl::InlinedVector<int64_t, 4UL>* strides,
+    absl::InlinedVector<xla::DExpr, 4UL>* begin_expr = nullptr,
+    absl::InlinedVector<xla::DExpr, 4UL>* end_expr = nullptr,
     StridedSliceShapeSpec* shape_spec = nullptr);
 
 // Simple class for determining if it is possible to broadcast a tensor to a
diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD
index e94208f90f4537..3988745801be74 100644
--- a/third_party/xla/xla/BUILD
+++ b/third_party/xla/xla/BUILD
@@ -460,6 +460,7 @@ cc_library(
         "layout_util.cc",
         "primitive_util.cc",
         "shape.cc",
+        "shape_expr.cc",
         "shape_partition.cc",
         "shape_util.cc",
     ],
@@ -469,6 +470,7 @@ cc_library(
         "layout_util.h",
         "primitive_util.h",
         "shape.h",
+        "shape_expr.h",
         "shape_partition.h",
         "shape_util.h",
     ],
diff --git a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc
index 2ce99c4b93a878..ffc2ed3f8baf73 100644
--- a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.cc
@@ -147,7 +147,7 @@ llvm::StructType* KernelCallFrameTy(llvm::LLVMContext& ctx) {
   llvm::PointerType* ptr = llvm::PointerType::getUnqual(ctx);
   llvm::IntegerType* i64 = llvm::IntegerType::getInt64Ty(ctx);
   return llvm::StructType::create("XLA_CPU_KernelCallFrame", ptr, ptr, i64,
-                                  ptr);
+                                  ptr, i64);
 }
 
 llvm::FunctionType* KernelFunctionTy(llvm::LLVMContext& ctx) {
@@ -316,6 +316,34 @@ auto KernelApiIrBuilder::EmitKernelPrototype(
   return EmitKernelPrototype(module, name, arguments, results);
 }
 
+llvm::Value* KernelApiIrBuilder::EmitGetBatchDim(llvm::IRBuilderBase& builder,
+                                                 llvm::Value* call_frame) {
+  llvm::LLVMContext& ctx = builder.getContext();
+  llvm::Type* ptr = llvm::PointerType::get(ctx, 0);
+  llvm::IntegerType* i64 = llvm::IntegerType::getInt64Ty(ctx);
+  llvm::Value* bdim_gep =
+      builder.CreateStructGEP(call_frame_ty_, call_frame, 4, "bdim_gep");
+  llvm::Value* bdim_value = builder.CreateLoad(i64, bdim_gep, "bdim_value");
+
+#if defined(PRINT_BATCHSIZE)
+  // Print batch size
+  llvm::Function* function = builder.GetInsertBlock()->getParent();
+  llvm::Module* module = function->getParent();
+  llvm::FunctionType* printfType = llvm::FunctionType::get(
+      builder.getInt32Ty(), llvm::PointerType::get(builder.getInt8Ty(), 0),
+      true);
+  llvm::Value* funcNameStr =
+      builder.CreateGlobalStringPtr(function->getName());
+  llvm::FunctionCallee printfFunc =
+      module->getOrInsertFunction("printf", printfType);
+  llvm::Value* formatStr =
+      builder.CreateGlobalStringPtr("Function: %s, Batch size is : %lld!\n");
+  builder.CreateCall(printfFunc, {formatStr, funcNameStr, bdim_value});
+#endif
+
+  return bdim_value;
+}
+
 auto KernelApiIrBuilder::EmitKernelPrototype(
     llvm::Module& module, absl::string_view name,
     absl::Span<const KernelParameter> arguments,
@@ -394,6 +422,8 @@ auto KernelApiIrBuilder::EmitKernelPrototype(
     ir_results.push_back(std::move(ir_result));
   }
 
+  EmitGetBatchDim(b, call_frame);
+
   // Return null pointer to signal success as we do not support error handling
   // in the compiled host kernel.
   llvm::BasicBlock* return_block =
@@ -503,7 +533,8 @@ llvm_ir::IrArray KernelApiIrBuilder::EmitKernelArgument(
   const llvm::DataLayout& data_layout = llvm_module->getDataLayout();
   int64_t pointer_size = data_layout.getTypeStoreSize(builder.getPtrTy());
   int64_t byte_size = ShapeUtil::ByteSizeOf(shape, pointer_size);
-  llvm_ir::SetDereferenceableMetadataForLoad(data, byte_size);
+  if (!shape.has_dynamic_expr())
+    llvm_ir::SetDereferenceableMetadataForLoad(data,byte_size);
 
   // All buffers pointers passed to host kernels are expected to be invariant
   // over the whole program. Note the metadata is attached only to loading
diff --git a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h
index 58ffc5bc8594e3..08af675fcb6c37 100644
--- a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h
+++ b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder.h
@@ -147,9 +147,13 @@ class KernelApiIrBuilder {
   llvm_ir::IrArray EmitKernelArgument(llvm::IRBuilderBase& builder,
                                       llvm::Value* call_frame, int64_t index,
                                       const Shape& shape);
+
   llvm::Function* EmitKernelFunction(llvm::Module& module,
                                      absl::string_view name);
 
+  llvm::Value* EmitGetBatchDim(llvm::IRBuilderBase& builder,
+                               llvm::Value* call_frame);
+
  private:
   llvm::LLVMContext& context_;
 
diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel.cc b/third_party/xla/xla/backends/cpu/runtime/kernel.cc
index ac1a5d7181ec52..668e10f0a0dd09 100644
--- a/third_party/xla/xla/backends/cpu/runtime/kernel.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/kernel.cc
@@ -64,7 +64,7 @@ template <bool thread_dim_x_only>
 class Kernel::ParallelTask {
  public:
   ParallelTask(XLA_CPU_Kernel* kernel, Kernel::ThreadDim thread_dims,
-               absl::Span<const XLA_CPU_KernelArg> args);
+               size_t batch_size, absl::Span<const XLA_CPU_KernelArg> args);
 
   // Invokes a host kernel for a given task index.
   absl::Status operator()(size_t task_index) const;
@@ -76,6 +76,7 @@ class Kernel::ParallelTask {
 
   XLA_CPU_Kernel* kernel_;
   XLA_CPU_KernelThreadDim thread_dims_;
+  size_t batch_size_;
   absl::InlinedVector<XLA_CPU_KernelArg, 8> args_;
 
   size_t num_tasks_;
@@ -88,11 +89,13 @@ class Kernel::ParallelTask {
 template <bool thread_dim_x_only>
 Kernel::ParallelTask<thread_dim_x_only>::ParallelTask(
     XLA_CPU_Kernel* kernel, Kernel::ThreadDim thread_dims,
+    size_t batch_size,
     absl::Span<const XLA_CPU_KernelArg> args)
     : kernel_(kernel),
       thread_dims_({thread_dims.x, thread_dims.y, thread_dims.z}),
       args_(args.begin(), args.end()),
       num_tasks_(thread_dims_.x * thread_dims_.y * thread_dims_.z),
+      batch_size_(batch_size),
       stride_z_(thread_dims_.y * thread_dims_.x),
       stride_y_(thread_dims_.x) {}
 
@@ -103,7 +106,7 @@ absl::Status Kernel::ParallelTask<thread_dim_x_only>::operator()(
 
   XLA_CPU_KernelThread kernel_thread = Delinearize(task_index);
   XLA_CPU_KernelCallFrame call_frame = {&thread_dims_, &kernel_thread,
-                                        args_.size(), args_.data()};
+                                        args_.size(), args_.data(), batch_size_};
 
   XLA_CPU_KernelError* error = (*kernel_)(&call_frame);
 
@@ -138,12 +141,12 @@ Kernel::Kernel(unsigned arity, XLA_CPU_Kernel* kernel)
       kernel_(function_->kernel()),
       arity_(arity) {}
 
-absl::Status Kernel::Launch(const ThreadDim& thread_dims,
+absl::Status Kernel::Launch(const ThreadDim& thread_dims, size_t batch_size,
                             absl::Span<const DeviceMemoryBase> buffers) const {
-  return Launch(thread_dims, ConvertBuffersToKernelArgs(buffers));
+  return Launch(thread_dims, batch_size, ConvertBuffersToKernelArgs(buffers));
 }
 
-absl::Status Kernel::Launch(const ThreadDim& thread_dims,
+absl::Status Kernel::Launch(const ThreadDim& thread_dims, size_t batch_size,
                             absl::Span<const XLA_CPU_KernelArg> args) const {
   XLA_CPU_KernelThreadDim kernel_thread_dims = {
       thread_dims.x,
@@ -156,8 +159,9 @@ absl::Status Kernel::Launch(const ThreadDim& thread_dims,
       for (uint64_t x = 0; x < thread_dims.x; ++x) {
         XLA_CPU_KernelThread kernel_thread = {x, y, z};
 
-        XLA_CPU_KernelCallFrame call_frame = {
-            &kernel_thread_dims, &kernel_thread, args.size(), args.data()};
+        XLA_CPU_KernelCallFrame call_frame = {&kernel_thread_dims,
+                                              &kernel_thread, args.size(),
+                                              args.data(), batch_size};
 
         XLA_CPU_KernelError* error = (*kernel_)(&call_frame);
 
@@ -172,20 +176,23 @@ absl::Status Kernel::Launch(const ThreadDim& thread_dims,
 }
 
 tsl::AsyncValueRef<LaunchEvent> Kernel::Launch(
-    const ThreadDim& thread_dims, absl::Span<const DeviceMemoryBase> buffers,
+    const ThreadDim& thread_dims, size_t batch_size,
+    absl::Span<const DeviceMemoryBase> buffers,
     const Eigen::ThreadPoolDevice* device) const {
-  return Launch(thread_dims, ConvertBuffersToKernelArgs(buffers), device);
+  return Launch(thread_dims, batch_size, ConvertBuffersToKernelArgs(buffers),
+                device);
 }
 
 tsl::AsyncValueRef<LaunchEvent> Kernel::Launch(
-    const ThreadDim& thread_dims, absl::Span<const XLA_CPU_KernelArg> args,
+    const ThreadDim& thread_dims, size_t batch_size,
+    absl::Span<const XLA_CPU_KernelArg> args,
     const Eigen::ThreadPoolDevice* device) const {
   size_t num_tasks = thread_dims.x * thread_dims.y * thread_dims.z;
   CHECK_GT(num_tasks, 0) << "Number of tasks must be positive";  // Crash Ok
 
   // Short-circuit launch with a single task and run it in the caller thread.
   if (ABSL_PREDICT_TRUE(num_tasks == 1)) {
-    absl::Status launched = Launch(thread_dims, args);
+    absl::Status launched = Launch(thread_dims, batch_size, args);
     return ABSL_PREDICT_TRUE(launched.ok())
                ? OkLaunchEvent()
                : tsl::MakeErrorAsyncValueRef(std::move(launched));
@@ -197,11 +204,13 @@ tsl::AsyncValueRef<LaunchEvent> Kernel::Launch(
                        std::numeric_limits<uint16_t>::max());
 
   if (ABSL_PREDICT_TRUE(thread_dims.y == 1 && thread_dims.z == 1)) {
-    return Worker::Parallelize(device, num_workers, num_tasks,
-                               ParallelTask<true>(kernel_, thread_dims, args));
+    return Worker::Parallelize(
+        device, num_workers, num_tasks,
+        ParallelTask<true>(kernel_, thread_dims, batch_size, args));
   } else {
-    return Worker::Parallelize(device, num_workers, num_tasks,
-                               ParallelTask<false>(kernel_, thread_dims, args));
+    return Worker::Parallelize(
+        device, num_workers, num_tasks,
+        ParallelTask<false>(kernel_, thread_dims, batch_size, args));
   }
 }
 
diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel.h b/third_party/xla/xla/backends/cpu/runtime/kernel.h
index 98010905322f67..d03657356bfc7e 100644
--- a/third_party/xla/xla/backends/cpu/runtime/kernel.h
+++ b/third_party/xla/xla/backends/cpu/runtime/kernel.h
@@ -73,13 +73,25 @@ class Kernel {
 
   // Calls the kernel once in the caller thread for a thread dim (0,0,0).
   // This is a fast path for small host kernels that have just one thread.
-  absl::Status CallOnce(absl::Span<const XLA_CPU_KernelArg> args) const;
+  absl::Status CallOnce(absl::Span<const XLA_CPU_KernelArg> args) const {
+    return CallOnce(args, 0);
+  }
+  absl::Status CallOnce(absl::Span<const XLA_CPU_KernelArg> args,
+                        size_t batch_size) const;
 
   // Launches the kernel on the current thread by iterating over all threads in
   // `thread_dims` and calling the kernel function.
   absl::Status Launch(const ThreadDim& thread_dims,
-                      absl::Span<const DeviceMemoryBase> buffers) const;
+                      absl::Span<const DeviceMemoryBase> buffers) const {
+    return Launch(thread_dims, 0, buffers);
+  }
   absl::Status Launch(const ThreadDim& thread_dims,
+                      absl::Span<const XLA_CPU_KernelArg> args) const {
+    return Launch(thread_dims, 0, args);
+  }
+  absl::Status Launch(const ThreadDim& thread_dims, size_t batch_size,
+                      absl::Span<const DeviceMemoryBase> buffers) const;
+  absl::Status Launch(const ThreadDim& thread_dims, size_t batch_size,
                       absl::Span<const XLA_CPU_KernelArg> args) const;
 
   // Launches the kernel by iterating over all threads in `thread_dims` and
@@ -90,9 +102,21 @@ class Kernel {
   // get the number of tasks that are expected to be completed.
   tsl::AsyncValueRef<LaunchEvent> Launch(
       const ThreadDim& thread_dims, absl::Span<const DeviceMemoryBase> buffers,
-      const Eigen::ThreadPoolDevice* device) const;
+      const Eigen::ThreadPoolDevice* device) const {
+    return Launch(thread_dims, 0, buffers, device);
+  }
   tsl::AsyncValueRef<LaunchEvent> Launch(
       const ThreadDim& thread_dims, absl::Span<const XLA_CPU_KernelArg> args,
+      const Eigen::ThreadPoolDevice* device) const {
+    return Launch(thread_dims, 0, args, device);
+  }
+  tsl::AsyncValueRef<LaunchEvent> Launch(
+      const ThreadDim& thread_dims, size_t batch_size,
+      absl::Span<const DeviceMemoryBase> buffers,
+      const Eigen::ThreadPoolDevice* device) const;
+  tsl::AsyncValueRef<LaunchEvent> Launch(
+      const ThreadDim& thread_dims, size_t batch_size,
+      absl::Span<const XLA_CPU_KernelArg> args,
       const Eigen::ThreadPoolDevice* device) const;
 
   // For host platform, we assume that a core is a thread, and we can run at
@@ -123,12 +147,13 @@ class Kernel {
 };
 
 inline ABSL_ATTRIBUTE_ALWAYS_INLINE absl::Status Kernel::CallOnce(
-    absl::Span<const XLA_CPU_KernelArg> args) const {
+    absl::Span<const XLA_CPU_KernelArg> args,
+    size_t batch_size) const {
   constexpr XLA_CPU_KernelThreadDim kernel_thread_dims = {1, 1, 1};
   constexpr XLA_CPU_KernelThread kernel_thread = {1, 1, 1};
 
   XLA_CPU_KernelCallFrame call_frame = {&kernel_thread_dims, &kernel_thread,
-                                        args.size(), args.data()};
+                                        args.size(), args.data(), batch_size};
 
   XLA_CPU_KernelError* error = (*kernel_)(&call_frame);
 
diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h b/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h
index cbe0568506385d..cf8016759ed217 100644
--- a/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h
+++ b/third_party/xla/xla/backends/cpu/runtime/kernel_c_api.h
@@ -72,6 +72,7 @@ typedef struct XLA_CPU_KernelCallFrame {
 
   size_t num_args;
   const XLA_CPU_KernelArg* args;
+  size_t batch_size;
 } XLA_CPU_KernelCallFrame;
 
 // Error reporting for host kernels. NULL means success.
diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc
index 90c15a09bd677d..7a1c0ec70c449b 100644
--- a/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.cc
@@ -223,7 +223,7 @@ KernelThunk<num_arguments, num_results>::ExecuteInternal(
 
   // Use a fast path if kernel called just once.
   if (ABSL_PREDICT_TRUE(call_once_)) {
-    TF_RETURN_IF_ERROR(kernel->CallOnce(kernel_args));
+    TF_RETURN_IF_ERROR(kernel->CallOnce(kernel_args, params.batch_size));
     return OkExecuteEvent();
   }
 
@@ -231,10 +231,12 @@ KernelThunk<num_arguments, num_results>::ExecuteInternal(
   // by scheduling tasks into it. HostKernel launch completion will
   // automatically signal KernelThunk execute completion.
   if (ABSL_PREDICT_TRUE(params.intra_op_threadpool)) {
-    return kernel->Launch(thread_dim_, kernel_args, params.intra_op_threadpool);
+    return kernel->Launch(thread_dim_, params.batch_size, kernel_args,
+                          params.intra_op_threadpool);
   }
 
-  TF_RETURN_IF_ERROR(kernel->Launch(thread_dim_, kernel_args));
+  TF_RETURN_IF_ERROR(
+      kernel->Launch(thread_dim_, params.batch_size, kernel_args));
   return OkExecuteEvent();
 }
 
diff --git a/third_party/xla/xla/backends/cpu/runtime/thunk.h b/third_party/xla/xla/backends/cpu/runtime/thunk.h
index d4a56ae88fa55b..5b7042ccb2803f 100644
--- a/third_party/xla/xla/backends/cpu/runtime/thunk.h
+++ b/third_party/xla/xla/backends/cpu/runtime/thunk.h
@@ -259,6 +259,7 @@ class Thunk {
     TaskRunner* task_runner = nullptr;
     CollectiveExecuteParams* collective_params = nullptr;
     CustomCallExecuteParams* custom_call_params = nullptr;
+    int64_t batch_size = 0;
     ExecuteSession session = ExecuteSession(ExecuteSession::kMaxWorkers,
                                             ExecuteSession::kSplitThreshold);
   };
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 33fa90f7e35e9e..961273151ed5e0 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -103,6 +103,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_cpu_use_fusion_emitters(true);
   opts.set_xla_cpu_use_thunk_runtime(true);
   opts.set_xla_cpu_use_xnnpack(false);
+  opts.set_xla_compile_batch_sizes("");
   opts.set_xla_cpu_experimental_xnn_graph_fusion_mode(
       DebugOptions::XNN_GRAPH_FUSION_MODE_DISABLED);
   opts.set_xla_cpu_parallel_codegen_split_count(32);
@@ -1003,6 +1004,15 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "`XNN_GRAPH_FUSION_MODE_DISABLED` - default value, "
       "`XNN_GRAPH_FUSION_MODE_GREEDY` - greedy extraction of "
       "XNNPACK-compatible subgraphs starting from root instructions."));
+  flag_list->push_back(tsl::Flag(
+      "xla_compile_batch_sizes",
+      string_setter_for(
+          &DebugOptions::set_xla_compile_batch_sizes),
+      debug_options->xla_compile_batch_sizes(),
+      "Comma-separated list of batch sizes to use for compilation, "
+      "use single value or start:end:step format. "
+      "e.g. 32, 64, 128, 10:100:10, "
+      "empty to use the nearest power of two."));
   flag_list->push_back(tsl::Flag(
       "xla_cpu_parallel_codegen_split_count",
       int32_setter_for(&DebugOptions::set_xla_cpu_parallel_codegen_split_count),
diff --git a/third_party/xla/xla/executable_run_options.h b/third_party/xla/xla/executable_run_options.h
index b377e91670efb9..38b13e5ec0726a 100644
--- a/third_party/xla/xla/executable_run_options.h
+++ b/third_party/xla/xla/executable_run_options.h
@@ -197,6 +197,13 @@ class ExecutableRunOptions {
 
   int32_t launch_id() const { return launch_id_; }
 
+  ExecutableRunOptions& set_batch_size(int64_t batch_size) {
+    batch_size_ = batch_size;
+    return *this;
+  }
+
+  int64_t batch_size() const { return batch_size_; }
+
   ExecutableRunOptions& set_run_id(RunId id);
   RunId run_id() const;
 
@@ -265,6 +272,7 @@ class ExecutableRunOptions {
   ExecutionProfile* execution_profile_ = nullptr;
   int rng_seed_ = 0;
   int32_t launch_id_ = 0;
+  int64_t batch_size_ = 0;
   stream_executor::Stream* device_to_host_stream_ = nullptr;
   stream_executor::Stream* host_to_device_stream_ = nullptr;
   ThenExecuteFunction* then_execute_function_ = nullptr;
diff --git a/third_party/xla/xla/hlo/builder/BUILD b/third_party/xla/xla/hlo/builder/BUILD
index 43dcce69c83579..0a508f54385cad 100644
--- a/third_party/xla/xla/hlo/builder/BUILD
+++ b/third_party/xla/xla/hlo/builder/BUILD
@@ -174,6 +174,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:protobuf",
         "@local_tsl//tsl/platform:stacktrace",
         "@local_tsl//tsl/platform:statusor",
     ],
diff --git a/third_party/xla/xla/hlo/builder/lib/approx_topk.cc b/third_party/xla/xla/hlo/builder/lib/approx_topk.cc
index 5b0f223302f4ed..9cb5fa5eefa2b8 100644
--- a/third_party/xla/xla/hlo/builder/lib/approx_topk.cc
+++ b/third_party/xla/xla/hlo/builder/lib/approx_topk.cc
@@ -131,10 +131,10 @@ XlaOp AggregateToTopKBuilder(XlaBuilder* builder,
                            reduction_computation, {reduction_dim});
     Shape op_shape = operands_shapes[0];
     op_shape.set_dimensions(reduction_dim, 1);
-    auto top1_vals =
-        Reshape(GetTupleElement(val_args, 0), op_shape.dimensions());
-    auto top1_args =
-        Reshape(GetTupleElement(val_args, 1), op_shape.dimensions());
+    auto top1_vals = Reshape(GetTupleElement(val_args, 0),
+                             op_shape.dimensions(), op_shape.expressions());
+    auto top1_args = Reshape(GetTupleElement(val_args, 1),
+                             op_shape.dimensions(), op_shape.expressions());
     return Tuple(builder, {top1_vals, top1_args});
   }
 
diff --git a/third_party/xla/xla/hlo/builder/lib/arithmetic.cc b/third_party/xla/xla/hlo/builder/lib/arithmetic.cc
index 84908846dae705..7bb8694351ce92 100644
--- a/third_party/xla/xla/hlo/builder/lib/arithmetic.cc
+++ b/third_party/xla/xla/hlo/builder/lib/arithmetic.cc
@@ -154,8 +154,9 @@ XlaOp ArgMinMax(XlaOp input, PrimitiveType output_type, int axis, bool is_min) {
     int64_t dimension_size = input_shape.dimensions(axis);
     auto index_type = dimension_size <= INT32_MAX ? S32 : output_type;
     XlaOp index_init_value = Zero(builder, index_type);
-    auto iota_shape =
-        ShapeUtil::MakeShape(index_type, input_shape.dimensions());
+    auto iota_shape = ShapeUtil::MakeShape(index_type, input_shape.dimensions(),
+                                           input_shape.expressions());
+
     XlaOp iota = Iota(builder, iota_shape, axis);
 
     XlaComputation reducer = CreateMinMaxComputation(
diff --git a/third_party/xla/xla/hlo/builder/lib/broadcast.cc b/third_party/xla/xla/hlo/builder/lib/broadcast.cc
index 1baec363b53a35..46f5410f2a817a 100644
--- a/third_party/xla/xla/hlo/builder/lib/broadcast.cc
+++ b/third_party/xla/xla/hlo/builder/lib/broadcast.cc
@@ -31,8 +31,9 @@ limitations under the License.
 
 namespace xla {
 
-absl::StatusOr<XlaOp> BroadcastTo(XlaOp input,
-                                  absl::Span<int64_t const> output_dims) {
+absl::StatusOr<XlaOp> BroadcastTo(
+    XlaOp input, absl::Span<int64_t const> output_dims,
+    absl::Span<const xla::DExpr> output_exprs) {
   XlaBuilder* builder = input.builder();
   TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input));
   absl::Span<int64_t const> input_dims = input_shape.dimensions();
@@ -48,6 +49,12 @@ absl::StatusOr<XlaOp> BroadcastTo(XlaOp input,
         absl::StrJoin(output_dims, ","), "]");
   }
 
+  if (!output_exprs.empty() && output_exprs.size() != output_dims.size()) {
+    return tsl::errors::InvalidArgument(
+        "output_exprs must be empty or have the same rank as output_dims: ",
+        output_exprs.size(), " vs ", output_dims.size());
+  }
+
   std::vector<int64_t> broadcast_dims;
   std::vector<int64_t> broadcast_shape;
   auto input_it = input_dims.rbegin();
@@ -79,15 +86,41 @@ absl::StatusOr<XlaOp> BroadcastTo(XlaOp input,
   }
   TF_RET_CHECK(input_it == input_dims.rend());
 
+  absl::Span<const DExpr> input_exprs = input_shape.expressions();
+  std::vector<DExpr> broadcast_exprs;
+  auto input_dim_et = input_dims.rbegin();
+  auto input_et = input_exprs.rbegin();
+  auto output_dim_et = output_dims.rbegin();
+  for (auto output_et = output_exprs.rbegin(); output_et != output_exprs.rend();
+       ++output_et, ++output_dim_et) {
+    if (input_et != input_exprs.rend()) {
+      if (*output_dim_et == *input_dim_et || *input_dim_et == 1 ||
+          **output_et == **input_et ||
+          (input_et->get()->is_constant() && input_et->get()->get_val() == 1)) {
+        broadcast_exprs.push_back(*output_et);
+      } else if (!(**output_et == **input_et)) {
+        broadcast_exprs.push_back(*input_et);
+        broadcast_exprs.push_back(*output_et / *input_et);
+      }
+      ++input_dim_et;
+      ++input_et;
+    } else {
+      broadcast_exprs.push_back(*output_et);
+    }
+  }
+
   absl::c_reverse(broadcast_dims);
   int broadcast_shape_size = broadcast_shape.size();
   for (int64_t& broadcast_dim : broadcast_dims) {
     broadcast_dim = broadcast_shape_size - broadcast_dim - 1;
   }
   absl::c_reverse(broadcast_shape);
-  XlaOp output = BroadcastInDim(input, broadcast_shape, broadcast_dims);
+  absl::c_reverse(broadcast_exprs);
+
+  XlaOp output =
+      BroadcastInDim(input, broadcast_shape, broadcast_dims, broadcast_exprs);
   if (broadcast_shape != output_dims) {
-    output = Reshape(output, output_dims);
+    output = Reshape(output, output_dims, output_exprs);
   }
   return output;
 }
diff --git a/third_party/xla/xla/hlo/builder/lib/broadcast.h b/third_party/xla/xla/hlo/builder/lib/broadcast.h
index 86cf39f64ddc82..dd5535f219b971 100644
--- a/third_party/xla/xla/hlo/builder/lib/broadcast.h
+++ b/third_party/xla/xla/hlo/builder/lib/broadcast.h
@@ -27,8 +27,9 @@ namespace xla {
 
 // Broadcasts 'input' up to shape 'output_dims', using TensorFlow broadcasting
 // rules. Supports broadcasting a dimension of size x to size x*y, i.e., tiling.
-absl::StatusOr<XlaOp> BroadcastTo(XlaOp input,
-                                  absl::Span<int64_t const> output_dims);
+absl::StatusOr<XlaOp> BroadcastTo(
+    XlaOp input, absl::Span<int64_t const> output_dims,
+    absl::Span<const xla::DExpr> output_exprs = {});
 
 }  // namespace xla
 
diff --git a/third_party/xla/xla/hlo/builder/lib/constants.cc b/third_party/xla/xla/hlo/builder/lib/constants.cc
index acfa2fe0b66e2c..da6a97f4088520 100644
--- a/third_party/xla/xla/hlo/builder/lib/constants.cc
+++ b/third_party/xla/xla/hlo/builder/lib/constants.cc
@@ -33,7 +33,8 @@ XlaOp Zero(XlaBuilder* builder, PrimitiveType type) {
 }
 
 XlaOp Zeros(XlaBuilder* builder, const Shape& shape) {
-  return Broadcast(Zero(builder, shape.element_type()), shape.dimensions());
+  return Broadcast(Zero(builder, shape.element_type()), shape.dimensions(),
+                   shape.expressions());
 }
 
 XlaOp ZerosLike(XlaOp prototype) {
diff --git a/third_party/xla/xla/hlo/builder/lib/matrix.cc b/third_party/xla/xla/hlo/builder/lib/matrix.cc
index e9fe29ea83ee6b..077a81c1ce38b4 100644
--- a/third_party/xla/xla/hlo/builder/lib/matrix.cc
+++ b/third_party/xla/xla/hlo/builder/lib/matrix.cc
@@ -209,7 +209,8 @@ XlaOp SetMatrixDiagonal(XlaOp matrix, XlaOp diag, int k) {
     }
 
     return Select(GetDiagonalMask(matrix, k),
-                  BroadcastInDim(diag, shape.dimensions(), broadcast_dims),
+                  BroadcastInDim(diag, shape.dimensions(), broadcast_dims,
+                                 shape.expressions()),
                   matrix);
   });
 }
@@ -341,18 +342,22 @@ xla::XlaOp EinsumInverseDiagonal(XlaOp x, absl::Span<const int64_t> config) {
     }
     TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x));
     std::vector<int64_t> broadcast_sizes;
+    std::vector<DExpr> broadcast_exprs;
     int64_t x_dim = 0;
     for (auto label = config.begin(); label != config.end(); ++label) {
       auto first_label = absl::c_find(config, *label);
       if (first_label == label) {
         broadcast_sizes.push_back(x_shape.dimensions(x_dim));
+        broadcast_exprs.push_back(x_shape.expressions(x_dim));
         ++x_dim;
       } else {
         broadcast_sizes.push_back(
             broadcast_sizes[first_label - config.begin()]);
+        broadcast_exprs.push_back(
+            broadcast_exprs[first_label - config.begin()]);
       }
     }
-    x = BroadcastInDim(x, broadcast_sizes, labels->at(2));
+    x = BroadcastInDim(x, broadcast_sizes, labels->at(2), broadcast_exprs);
     return EinsumDiagonalMask(x, config);
   });
 }
@@ -568,16 +573,20 @@ xla::XlaOp Einsum(xla::XlaOp x, absl::Span<const int64_t> x_config,
 
     int64_t dot_dim = 0;
     std::vector<int64_t> new_dims;
+    std::vector<DExpr> new_exprs;
     new_dims.reserve(output_rank);
+    new_exprs.reserve(output_rank);
     TF_ASSIGN_OR_RETURN(Shape dot_shape, builder->GetShape(dot));
     for (auto d : output_config) {
       if (is_output_only(d)) {
         new_dims.push_back(1);
+        new_exprs.push_back(DExpr::Const(1));
       } else {
         new_dims.push_back(dot_shape.dimensions(dot_dim));
+        new_exprs.push_back(dot_shape.expressions(dot_dim));
       }
     }
-    return Reshape(dot, new_dims);
+    return Reshape(dot, new_dims, new_exprs);
   });
 }
 
diff --git a/third_party/xla/xla/hlo/builder/lib/prng.cc b/third_party/xla/xla/hlo/builder/lib/prng.cc
index 661981f19c17ea..7bb1c619f824d5 100644
--- a/third_party/xla/xla/hlo/builder/lib/prng.cc
+++ b/third_party/xla/xla/hlo/builder/lib/prng.cc
@@ -41,8 +41,9 @@ namespace xla {
 xla::XlaOp ConcatScalars(xla::XlaBuilder* builder,
                          absl::Span<const xla::XlaOp> scalars) {
   std::vector<xla::XlaOp> vectors;
-  absl::c_transform(scalars, std::back_inserter(vectors),
-                    [](xla::XlaOp x) { return xla::Reshape(x, {1}); });
+  absl::c_transform(scalars, std::back_inserter(vectors), [](xla::XlaOp x) {
+    return xla::Reshape(x, {1}, {xla::DExpr::Const(1)});
+  });
   return ConcatInDim(builder, vectors, 0);
 }
 
@@ -154,7 +155,8 @@ std::pair<ThreeFry2x32State, XlaOp> GetThreeFryInputsAndUpdatedState(
   XlaBuilder* builder = initial_state.builder();
   auto u64_shape = ShapeUtil::MakeShape(U64, shape.dimensions());
   // initial_state is an R1, so reshape it to a scalar.
-  auto input_u64 = Broadcast(Reshape(initial_state, {}), shape.dimensions());
+  auto input_u64 = Broadcast(Reshape(initial_state, {}), shape.dimensions(),
+                             shape.expressions());
   int64_t trailing_dims_product = 1;
   for (int64_t i = shape.dimensions().size() - 1; i >= 0; --i) {
     if (shape.dimensions(i) < 2) {
@@ -245,8 +247,12 @@ XlaOp CombineShapePair(absl::Span<const XlaOp> pair,
       original_shape.dimensions(shape_pair.split_dim);
   std::vector<int64_t> reshape_dims(original_shape.dimensions().begin(),
                                     original_shape.dimensions().end());
+  std::vector<DExpr> reshape_exprs(original_shape.expressions().begin(),
+                                   original_shape.expressions().end());
   reshape_dims[shape_pair.split_dim] = RoundUpTo<int64_t>(pre_split_size, 2);
-  result = Reshape(result, reshape_dims);
+  reshape_exprs[shape_pair.split_dim] =
+      DExpr::Const(RoundUpTo<int64_t>(pre_split_size, 2));
+  result = Reshape(result, reshape_dims, reshape_exprs);
   if (reshape_dims[shape_pair.split_dim] != pre_split_size) {
     result = Slice(result,
                    std::vector<int64_t>(original_shape.dimensions().size(), 0),
@@ -453,11 +459,11 @@ RngOutput PhiloxRngBit32(XlaOp op_key, XlaOp initial_state,
   }
   XlaOp numbers = ConcatInDim(builder, {bits[0], bits[1], bits[2], bits[3]},
                               /*dimension=*/1);
-  numbers = Reshape(numbers, {bits_len * 4});
+  numbers = Reshape(numbers, {bits_len * 4}, {});
   numbers = Slice(numbers, /*start_indices=*/{0},
                   /*limit_indices=*/{num_elems},
                   /*strides=*/{1});
-  return {Reshape(numbers, shape.dimensions()), new_state};
+  return {Reshape(numbers, shape.dimensions(), shape.expressions()), new_state};
 }
 
 // Generates an array of primitive type U16 with the given shape containing
@@ -507,7 +513,7 @@ RngOutput PhiloxRngBit64(XlaOp op_key, XlaOp initial_state,
   numbers = Slice(numbers, /*start_indices=*/{0},
                   /*limit_indices=*/{num_elems},
                   /*strides=*/{1});
-  return {Reshape(numbers, shape.dimensions()), new_state};
+  return {Reshape(numbers, shape.dimensions(), shape.expressions()), new_state};
 }
 
 XlaOp ConvertRandomBitsToUniformFloatingPoint(XlaOp bits, XlaOp minval,
diff --git a/third_party/xla/xla/hlo/builder/lib/slicing.cc b/third_party/xla/xla/hlo/builder/lib/slicing.cc
index ae0f6b987497a1..553936d5790a95 100644
--- a/third_party/xla/xla/hlo/builder/lib/slicing.cc
+++ b/third_party/xla/xla/hlo/builder/lib/slicing.cc
@@ -168,13 +168,16 @@ XlaOp TorchGather(XlaOp input, XlaOp index, int64_t dim, bool sparse) {
       std::vector<int64_t> index_broadcast_dims;
       std::vector<int64_t> input_broadcast_dims;
       std::vector<int64_t> sizes;
+      std::vector<DExpr> expressions;
       sizes.reserve(index_shape.dimensions().size());
+      expressions.reserve(index_shape.expressions().size());
       for (int64_t i = 0; i < index_shape.dimensions().size(); ++i) {
         if (i < dim) {
           input_broadcast_dims.push_back(i);
           index_broadcast_dims.push_back(i);
         } else if (i == dim) {
           sizes.push_back(input_shape.dimensions(i));
+          expressions.push_back(input_shape.expressions(i));
           input_broadcast_dims.push_back(i);
           index_broadcast_dims.push_back(i + 1);
         } else {
@@ -182,15 +185,18 @@ XlaOp TorchGather(XlaOp input, XlaOp index, int64_t dim, bool sparse) {
           index_broadcast_dims.push_back(i + 1);
         }
         sizes.push_back(index_shape.dimensions(i));
+        expressions.push_back(index_shape.expressions(i));
       }
-      auto mask = Eq(
-          BroadcastInDim(index, sizes, index_broadcast_dims),
-          Iota(builder, ShapeUtil::MakeShape(index_shape.element_type(), sizes),
-               dim));
+      auto mask =
+          Eq(BroadcastInDim(index, sizes, index_broadcast_dims, expressions),
+             Iota(builder,
+                  ShapeUtil::MakeShape(index_shape.element_type(), sizes,
+                                       expressions),
+                  dim));
       auto masked_input = Select(
-          mask, BroadcastInDim(input, sizes, input_broadcast_dims),
-          Zeros(builder,
-                ShapeUtil::MakeShape(input_shape.element_type(), sizes)));
+          mask, BroadcastInDim(input, sizes, input_broadcast_dims, expressions),
+          Zeros(builder, ShapeUtil::MakeShape(input_shape.element_type(), sizes,
+                                              expressions)));
       return Reduce(masked_input, Zero(builder, input_shape.element_type()),
                     CreateScalarIdentityWithZeroComputation(
                         input_shape.element_type(), builder),
@@ -203,7 +209,8 @@ XlaOp TorchGather(XlaOp input, XlaOp index, int64_t dim, bool sparse) {
     to_concat.reserve(input_shape.dimensions().size());
     for (int64_t i = 0; i < input_shape.dimensions().size(); ++i) {
       if (i == dim) {
-        to_concat.push_back(Reshape(index, index_shape.dimensions()));
+        to_concat.push_back(Reshape(index, index_shape.dimensions(),
+                                    index_shape.expressions()));
       } else {
         to_concat.push_back(Iota(builder, index_shape, i));
       }
@@ -229,27 +236,33 @@ XlaOp TorchScatterDense(XlaOp input, XlaOp index, XlaOp src, int64_t dim,
     TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input));
     std::vector<int64_t> index_broadcast_dims;
     std::vector<int64_t> sizes;
+      std::vector<DExpr> expressions;
     const auto rank = index_shape.dimensions().size();
     sizes.reserve(rank + 1);
+    expressions.reserve(rank + 1);
     for (int64_t i = 0; i < index_shape.dimensions().size(); ++i) {
       if (i < dim) {
         index_broadcast_dims.push_back(i);
       } else {
         if (i == dim) {
           sizes.push_back(input_shape.dimensions(i));
+          expressions.push_back(input_shape.expressions(i));
         }
         index_broadcast_dims.push_back(i + 1);
       }
       sizes.push_back(index_shape.dimensions(i));
+      expressions.push_back(index_shape.expressions(i));
     }
     auto mask =
-        Eq(BroadcastInDim(index, sizes, index_broadcast_dims),
+        Eq(BroadcastInDim(index, sizes, index_broadcast_dims, expressions),
            Iota(builder,
-                ShapeUtil::MakeShape(index_shape.element_type(), sizes), dim));
-    auto masked_src =
-        Select(mask, BroadcastInDim(src, sizes, index_broadcast_dims),
-               Zeros(builder,
-                     ShapeUtil::MakeShape(input_shape.element_type(), sizes)));
+                ShapeUtil::MakeShape(index_shape.element_type(), sizes,
+                                     expressions),
+                dim));
+    auto masked_src = Select(
+        mask, BroadcastInDim(src, sizes, index_broadcast_dims, expressions),
+        Zeros(builder, ShapeUtil::MakeShape(input_shape.element_type(), sizes,
+                                            expressions)));
 
     return combiner(
         input,
@@ -287,7 +300,8 @@ XlaOp TorchIndexSelect(XlaOp input, XlaOp index, int64_t dim,
       for (int64_t batch_dim = 0; batch_dim < batch_dims; ++batch_dim) {
         to_concat.push_back(Iota(builder, iota_shape, batch_dim));
       }
-      to_concat.push_back(Reshape(index, index_shape.dimensions()));
+      to_concat.push_back(
+          Reshape(index, index_shape.dimensions(), index_shape.expressions()));
       index = ConcatInDim(builder, to_concat, gather_dnums.index_vector_dim());
     }
     for (int64_t i = 0; i < input_shape.dimensions().size(); ++i) {
diff --git a/third_party/xla/xla/hlo/builder/lib/svd.cc b/third_party/xla/xla/hlo/builder/lib/svd.cc
index 561c107ba8085a..9254085e69a85e 100644
--- a/third_party/xla/xla/hlo/builder/lib/svd.cc
+++ b/third_party/xla/xla/hlo/builder/lib/svd.cc
@@ -152,9 +152,9 @@ absl::StatusOr<HouseHolderResult> HouseRow(
   auto beta = Div(ScalarLike(v_0j, 2.0),
                   (Square(Div(sigma, v_0j, broadcast_dims)) + one));
 
-  v = Select(
-      BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v,
-      v / v_0j);
+  v = Select(BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(),
+                            broadcast_dims, x_shape.expressions()),
+             v, v / v_0j);
   v = Select(Eq(idx, j), zeros + one, v);
 
   beta = Select(Lt(Add(sigma, ZerosLike(beta), broadcast_dims), eps),
@@ -219,9 +219,9 @@ absl::StatusOr<HouseHolderResult> HouseCol(
   auto beta = Div(ScalarLike(v_0i, 2.0),
                   (Square(Div(sigma, v_0i, broadcast_dims)) + one));
 
-  v = Select(
-      BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(), broadcast_dims), v,
-      v / v_0i);
+  v = Select(BroadcastInDim(Lt(sigma, eps), x_shape.dimensions(),
+                            broadcast_dims, x_shape.expressions()),
+             v, v / v_0i);
   v = Select(Eq(idx, i), zeros + one, v);
 
   beta = Select(Lt(Add(sigma, ZerosLike(beta), broadcast_dims), eps),
@@ -582,11 +582,11 @@ absl::StatusOr<XlaOp> ComputeToleranceComparison(XlaOp w, XlaOp epsilon) {
   diag = Select(Lt(diag, ZerosLike(diag)), -diag, diag);
   std::vector<int64_t> broadcasted_dims(num_dims - 1);
   std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0);
-  auto broadcast_to_rows =
-      BroadcastInDim(diag, shape.dimensions(), broadcasted_dims);
+  auto broadcast_to_rows = BroadcastInDim(
+      diag, shape.dimensions(), broadcasted_dims, shape.expressions());
   broadcasted_dims.back() = num_dims - 1;
-  auto broadcast_to_columns =
-      BroadcastInDim(diag, shape.dimensions(), broadcasted_dims);
+  auto broadcast_to_columns = BroadcastInDim(
+      diag, shape.dimensions(), broadcasted_dims, shape.expressions());
   // Compute tolerance = w_{i,i} * w_{j,j} * epsilon^2
   // Use at least F32 precision to avoid precision issues with small denormal.
   XlaOp tolerance;
@@ -745,6 +745,7 @@ absl::StatusOr<SVDResult> SortBySingularValuesAndPostProcessing(
   TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(result.d));
   const int64_t num_dims = shape.dimensions().size();
   auto dimensions = shape.dimensions();
+  auto expressions = shape.expressions();
   const int64_t m = ShapeUtil::GetDimension(shape, -2);
   const int64_t n = ShapeUtil::GetDimension(shape, -1);
 
@@ -763,7 +764,7 @@ absl::StatusOr<SVDResult> SortBySingularValuesAndPostProcessing(
   d = Select(Ge(d, zeros), d, -d);
   result.v = Mul(result.v, sign, broadcast_dims);
 
-  d = BroadcastInDim(d, dimensions, broadcast_dims);
+  d = BroadcastInDim(d, dimensions, broadcast_dims, expressions);
 
   // As m >= n, only first n column vectors need to be permuted, and the rest of
   // m - n vectors are appended after the sorting is done.
diff --git a/third_party/xla/xla/hlo/builder/value_inference.cc b/third_party/xla/xla/hlo/builder/value_inference.cc
index 169e8284aa96f5..ba72ae387e7b00 100644
--- a/third_party/xla/xla/hlo/builder/value_inference.cc
+++ b/third_party/xla/xla/hlo/builder/value_inference.cc
@@ -795,6 +795,26 @@ absl::StatusOr<PostorderDFSNode> PostorderDFSVisitor::AnalyzeUpperBound(
             return Literal::CreateFromProto(root->literal());
           }
         });
+      } else if (root->custom_call_target() == "GetExpressionValue") {
+        return PostorderDFSNode()
+            .AddDependency(root->operand_ids(0),
+                           PostorderDFSNodeType::kConstantUpperBound,
+                           InferenceContext({}, {}))
+            .AddVisit([](Literal carrier) -> absl::StatusOr<Literal> {
+              if (carrier.shape().dimensions_size() != 1 ||
+                  carrier.shape().dimensions(0) != 1) {
+                return InvalidArgument(
+                    "GetExpressionValue carrier must be rank-1 with one "
+                    "element, got %s",
+                    carrier.shape().ToString());
+              }
+              if (carrier.shape().element_type() != S32) {
+                return InvalidArgument(
+                    "GetExpressionValue carrier must be s32, got %s",
+                    carrier.shape().ToString());
+              }
+              return LiteralUtil::CreateR0<int32_t>(carrier.Get<int32_t>({0}));
+            });
       } else if (root->custom_call_target() == "Sharding") {
         return PostorderDFSNode()
             .AddDependency(root->operand_ids(0),
@@ -975,6 +995,14 @@ absl::StatusOr<PostorderDFSNode> PostorderDFSVisitor::AnalyzeConstant(
             return Literal::CreateFromProto(root->literal());
           }
         });
+      } else if (root->custom_call_target() == "GetExpressionValue") {
+        return PostorderDFSNode().AddVisit(
+            [root, context](absl::Span<Literal>) -> absl::StatusOr<Literal> {
+              TF_ASSIGN_OR_RETURN(Shape root_shape,
+                                  Shape::FromProto(root->shape()));
+              return CreateGarbageLiteral(
+                  ShapeUtil::GetSubshape(root_shape, context.shape_index));
+            });
       } else if (root->custom_call_target() == "Sharding") {
         return PostorderDFSNode()
             .AddDependency(root->operand_ids(0),
@@ -1507,6 +1535,16 @@ absl::StatusOr<PostorderDFSNode> PostorderDFSVisitor::AnalyzeIsDynamic(
             }
           }
         });
+      } else if (root->custom_call_target() == "GetExpressionValue") {
+        return PostorderDFSNode().AddVisit(
+            [type, root]() -> absl::StatusOr<Literal> {
+              TF_ASSIGN_OR_RETURN(Shape root_shape,
+                                  Shape::FromProto(root->shape()));
+              if (type == PostorderDFSNodeType::kValueIsDynamic) {
+                return CreatePredLiteral(true, root_shape);
+              }
+              return CreatePredLiteral(false, root_shape);
+            });
       } else if (root->custom_call_target() == "Sharding") {
         return result.AddVisit([](Literal operand) { return operand; });
       } else {
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc
index 93d7782de50e03..7c868a41610406 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder.cc
+++ b/third_party/xla/xla/hlo/builder/xla_builder.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
+#include "tsl/platform/protobuf.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/str_split.h"
@@ -56,6 +57,7 @@ limitations under the License.
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/permutation_util.h"
+#include "xla/printer.h"
 #include "xla/primitive_util.h"
 #include "xla/service/hlo.pb.h"
 #include "xla/service/shape_inference.h"
@@ -518,6 +520,20 @@ static std::string ShapeToString(const ShapeProto& shape) {
   return absl::StrCat("[", absl::StrJoin(shape.dimensions(), ", "), "]");
 }
 
+static std::vector<ExpressionProto> ContentsToProto(
+    absl::Span<const DExpr> contents) {
+  std::vector<ExpressionProto> protos;
+  protos.reserve(contents.size());
+  for (const DExpr& expr : contents) {
+    ExpressionProto proto;
+    if (expr) {
+      expr.to_proto(&proto);
+    }
+    protos.push_back(std::move(proto));
+  }
+  return protos;
+}
+
 void XlaBuilder::ToStringHelper(std::string* out, int ident,
                                 int64_t op_handle) const {
   const HloInstructionProto& instr =
@@ -707,6 +723,37 @@ absl::Status XlaBuilder::SetInstructionFrontendAttribute(const XlaOp op,
   return absl::OkStatus();
 }
 
+absl::Status XlaBuilder::SetInstructionContents(XlaOp op,
+                                                std::vector<DExpr> contents) {
+  auto it = handle_to_index_.find(op.handle());
+  if (it == handle_to_index_.end()) {
+    return InvalidArgument("No XlaOp with handle %d", op.handle());
+  }
+  const bool has_dynamic_content =
+      absl::c_any_of(contents, [](const DExpr& expr) {
+        return expr && expr->is_dynamic();
+      });
+  if (!has_dynamic_content) {
+    VLOG(1) << "SetInstructionContents clearing non-dynamic contents for op "
+              << op.handle() << " size=" << contents.size();
+    contents.clear();
+  } else {
+    VLOG(1) << "SetInstructionContents keeping dynamic contents for op "
+              << op.handle() << " size=" << contents.size();
+  }
+  instruction_contents_.at(it->second) = std::move(contents);
+  return absl::OkStatus();
+}
+
+absl::StatusOr<const std::vector<DExpr>*> XlaBuilder::GetInstructionContents(
+    XlaOp op) const {
+  auto it = handle_to_index_.find(op.handle());
+  if (it == handle_to_index_.end()) {
+    return InvalidArgument("No XlaOp with handle %d", op.handle());
+  }
+  return &instruction_contents_.at(it->second);
+}
+
 absl::Status XlaBuilder::SetInstructionSharding(
     XlaOp op, const std::optional<OpSharding>& sharding) {
   TF_ASSIGN_OR_RETURN(auto instr_proto, LookUpMutableInstruction(op));
@@ -784,7 +831,13 @@ absl::StatusOr<XlaComputation> XlaBuilder::Build(
   *entry.mutable_program_shape() = program_shape.ToProto();
   entry.set_root_id(root_id);
 
-  for (auto& instruction : instructions_) {
+  for (size_t index = 0; index < instructions_.size(); ++index) {
+    auto& instruction = instructions_[index];
+    if (!instruction_contents_[index].empty()) {
+      for (const auto& content : ContentsToProto(instruction_contents_[index])) {
+        *instruction.add_contents() = content;
+      }
+    }
     // Ensures that the instruction names are unique among the whole graph.
     instruction.set_name(
         GetFullName(instruction.name(), kNameSeparator, instruction.id()));
@@ -810,6 +863,7 @@ absl::StatusOr<XlaComputation> XlaBuilder::Build(
   // Clear data held by this builder.
   this->instructions_.clear();
   this->instruction_shapes_.clear();
+  this->instruction_contents_.clear();
   this->handle_to_index_.clear();
   this->embedded_.clear();
   this->parameter_numbers_.clear();
@@ -1017,22 +1071,24 @@ absl::StatusOr<XlaOp> XlaBuilder::AddBroadcastSequence(
   Shape broadcast_shape =
       ShapeUtil::ChangeElementType(output_shape, operand_shape->element_type());
 
-  // Do explicit broadcast for scalar.
-  if (ShapeUtil::IsScalar(*operand_shape)) {
-    return InDimBroadcast(ShapeUtil::MakeStaticShape(broadcast_shape), operand,
-                          {});
-  }
+    // Do explicit broadcast for scalar.
+    if (ShapeUtil::IsScalar(*operand_shape)) {
+      return InDimBroadcast(ShapeUtil::MakeStaticShape(broadcast_shape),
+                            operand, {});
+    }
 
   // Do explicit broadcast for degenerate broadcast.
   std::vector<int64_t> broadcast_dimensions;
   std::vector<int64_t> reshaped_dimensions;
   std::vector<bool> reshaped_dynamic_dimensions;
+  std::vector<DExpr> reshaped_expressions;
   for (int i = 0; i < operand_shape->dimensions().size(); i++) {
     if (operand_shape->dimensions(i) == output_shape.dimensions(i)) {
       broadcast_dimensions.push_back(i);
       reshaped_dimensions.push_back(operand_shape->dimensions(i));
       reshaped_dynamic_dimensions.push_back(
           operand_shape->is_dynamic_dimension(i));
+      reshaped_expressions.push_back(operand_shape->expressions(i));
     } else {
       TF_RET_CHECK(operand_shape->dimensions(i) == 1 &&
                    operand_shape->is_static_dimension(i))
@@ -1046,7 +1102,8 @@ absl::StatusOr<XlaOp> XlaBuilder::AddBroadcastSequence(
 
   Shape reshaped_shape =
       ShapeUtil::MakeShape(operand_shape->element_type(), reshaped_dimensions,
-                           reshaped_dynamic_dimensions);
+                           reshaped_dynamic_dimensions,
+                           reshaped_expressions);
 
   // Eliminate the size one dimensions.
   // The added reshape reduces the rank of the tensor. Hence we cannot directly
@@ -1094,15 +1151,20 @@ absl::StatusOr<XlaOp> BroadcastToTargetRank(
     return origin;
   }
 
-  // Update target_size with origin sizes using broadcast_dimensions
+  // Update target_size and target_exp with origin sizes and expressions using
+  // broadcast_dimensions
   absl::Span<const int64_t> target_dimensions = target_shape.dimensions();
+  absl::Span<const DExpr> target_expressions = target_shape.expressions();
   std::vector<int64_t> target_size{target_dimensions.begin(),
                                    target_dimensions.end()};
+  std::vector<DExpr> target_exp(target_expressions.begin(),
+                                target_expressions.end());
   for (int64_t origin_dim = 0; origin_dim < origin_rank; origin_dim++) {
     int64_t target_dim = broadcast_dimensions[origin_dim];
     target_size[target_dim] = origin_shape.dimensions(origin_dim);
+    target_exp[target_dim] = origin_shape.expressions(origin_dim);
   }
-  return BroadcastInDim(origin, target_size, broadcast_dimensions);
+  return BroadcastInDim(origin, target_size, broadcast_dimensions, target_exp);
 }
 
 // Extract the `num_dims` counts of dimension sizes from the `op`. First,
@@ -1120,7 +1182,7 @@ absl::StatusOr<std::vector<XlaOp>> ExtractDimensionSizesAndPadOnesToLeft(
             ? ConstantR1<int32_t>(
                   /*builder=*/builder,
                   /*values=*/{static_cast<int32_t>(op_shape->dimensions(i))})
-            : Reshape(GetDimensionSize(op, i), {1}));
+            : Reshape(GetDimensionSize(op, i), {1}, {xla::DExpr::Const(1)}));
   }
   return op_dims;
 }
@@ -1142,7 +1204,8 @@ absl::StatusOr<XlaOp> BroadcastScalarToOutputShapeWithUnbounded(
             ? ConstantR1<int32_t>(
                   /*builder=*/builder,
                   /*values=*/{static_cast<int32_t>(output_shape.dimensions(i))})
-            : Reshape(GetDimensionSize(output, i), {1});
+            : Reshape(GetDimensionSize(output, i), {1},
+                      {xla::DExpr::Const(1)});
   }
   return MhloDynamicBroadcastInDim(
       scalar, /*output_dimensions=*/ConcatInDim(builder, output_sizes, 0), {},
@@ -1525,20 +1588,13 @@ XlaOp XlaBuilder::Parameter(
 }
 
 XlaOp XlaBuilder::Broadcast(XlaOp operand,
-                            absl::Span<const int64_t> broadcast_sizes) {
+                            absl::Span<const int64_t> broadcast_sizes,
+                            absl::Span<const DExpr> broadcast_exprs) {
   return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
-    TF_ASSIGN_OR_RETURN(
-        const Shape& shape,
-        ShapeInference::InferBroadcastShape(*operand_shape, broadcast_sizes));
-
-    // The client-level broadcast op just appends dimensions on the left (adds
-    // lowest numbered dimensions). The HLO broadcast instruction is more
-    // flexible and can add new dimensions anywhere. The instruction's
-    // dimensions field maps operand dimensions to dimensions in the broadcast
-    // output, so to append dimensions on the left the instruction's dimensions
-    // should just be the n highest dimension numbers of the output shape where
-    // n is the number of input dimensions.
+    TF_ASSIGN_OR_RETURN(const Shape& shape,
+                        ShapeInference::InferBroadcastShape(
+                            *operand_shape, broadcast_sizes, broadcast_exprs));
     const int64_t operand_rank = operand_shape->dimensions().size();
     std::vector<int64_t> dimensions(operand_rank);
     for (int i = 0; i < operand_rank; ++i) {
@@ -1550,14 +1606,13 @@ XlaOp XlaBuilder::Broadcast(XlaOp operand,
 
 XlaOp XlaBuilder::BroadcastInDim(
     XlaOp operand, absl::Span<const int64_t> out_dim_size,
-    absl::Span<const int64_t> broadcast_dimensions) {
+    absl::Span<const int64_t> broadcast_dimensions,
+    absl::Span<const DExpr> out_dim_exp) {
   return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
-    // Output shape, in the case of degenerate broadcast, the out_dim_size is
-    // not necessarily the same as the dimension sizes of the output shape.
-    TF_ASSIGN_OR_RETURN(auto output_shape,
-                        ShapeUtil::MakeValidatedShape(
-                            operand_shape->element_type(), out_dim_size));
+    TF_ASSIGN_OR_RETURN(auto output_shape, ShapeUtil::MakeValidatedShape(
+                                               operand_shape->element_type(),
+                                               out_dim_size, out_dim_exp));
     TF_RET_CHECK(!output_shape.is_unbounded_dynamic())
         << "BroadcastInDim output must shape be static or bounded dynamic "
         << ShapeUtil::HumanString(output_shape);
@@ -1584,6 +1639,15 @@ XlaOp XlaBuilder::BroadcastInDim(
                            .status());
     std::vector<int64_t> in_dim_size(out_dim_size.begin(), out_dim_size.end());
     std::vector<bool> in_dim_dynamic(out_dim_size.size(), false);
+    std::vector<DExpr> in_expressions(out_dim_exp.begin(), out_dim_exp.end());
+
+    if (out_dim_exp.empty()) {
+      in_expressions.reserve(out_dim_size.size());
+      std::transform(out_dim_size.begin(), out_dim_size.end(),
+                     std::back_inserter(in_expressions),
+                     [](int d) { return DExpr::Const(d); });
+    }
+
     for (int i = 0; i < broadcast_rank; i++) {
       in_dim_size[broadcast_dimensions[i]] =
           (operand_shape->is_unbounded_dynamic_dimension(i))
@@ -1591,19 +1655,17 @@ XlaOp XlaBuilder::BroadcastInDim(
               : operand_shape->dimensions(i);
       in_dim_dynamic[broadcast_dimensions[i]] =
           operand_shape->is_bounded_dynamic_dimension(i);
+      in_expressions[broadcast_dimensions[i]] = operand_shape->expressions(i);
     }
-    const auto& in_dim_shape = ShapeUtil::MakeShape(
-        operand_shape->element_type(), in_dim_size, in_dim_dynamic);
+    const auto& in_dim_shape =
+        ShapeUtil::MakeShape(operand_shape->element_type(), in_dim_size,
+                             in_dim_dynamic, in_expressions);
     TF_ASSIGN_OR_RETURN(
         XlaOp in_dim_broadcast,
         InDimBroadcast(in_dim_shape, operand, broadcast_dimensions));
-
-    // If broadcast is not degenerate, return broadcasted result.
     if (ShapeUtil::Equal(in_dim_shape, output_shape)) {
       return in_dim_broadcast;
     }
-
-    // Otherwise handle degenerate broadcast case.
     return AddBroadcastSequence(output_shape, in_dim_broadcast);
   });
 }
@@ -1637,10 +1699,46 @@ XlaOp XlaBuilder::Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
   });
 }
 
+XlaOp XlaBuilder::Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
+                        absl::Span<const int64_t> limit_indices,
+                        absl::Span<const DExpr> start_exprs,
+                        absl::Span<const DExpr> limit_exprs,
+                        absl::Span<const int64_t> strides) {
+  return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
+    TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
+    TF_ASSIGN_OR_RETURN(
+        Shape shape, ShapeInference::InferSliceShape(*operand_shape,
+                                                     start_indices,
+                                                     limit_indices, strides,
+                                                     start_exprs, limit_exprs));
+    return SliceInternal(shape, operand, start_indices, limit_indices,
+                         start_exprs, limit_exprs, strides);
+  });
+}
+
+absl::StatusOr<XlaOp> XlaBuilder::SliceInternal(
+    const Shape& shape, XlaOp operand, absl::Span<const int64_t> start_indices,
+    absl::Span<const int64_t> limit_indices,
+    absl::Span<const int64_t> strides) {
+  HloInstructionProto instr;
+  *instr.mutable_shape() = shape.ToProto();
+  for (int i = 0, end = start_indices.size(); i < end; i++) {
+    auto* slice_config = instr.add_slice_dimensions();
+    slice_config->set_start(start_indices[i]);
+    slice_config->set_limit(limit_indices[i]);
+    slice_config->set_stride(strides[i]);
+  }
+  return AddInstruction(std::move(instr), HloOpcode::kSlice, {operand});
+}
+
 absl::StatusOr<XlaOp> XlaBuilder::SliceInternal(
     const Shape& shape, XlaOp operand, absl::Span<const int64_t> start_indices,
     absl::Span<const int64_t> limit_indices,
+    absl::Span<const DExpr> start_exprs,
+    absl::Span<const DExpr> limit_exprs,
     absl::Span<const int64_t> strides) {
+  CHECK(start_exprs.empty() || start_exprs.size() == start_indices.size());
+  CHECK(limit_exprs.empty() || limit_exprs.size() == start_indices.size());
   HloInstructionProto instr;
   *instr.mutable_shape() = shape.ToProto();
   for (int i = 0, end = start_indices.size(); i < end; i++) {
@@ -1648,6 +1746,14 @@ absl::StatusOr<XlaOp> XlaBuilder::SliceInternal(
     slice_config->set_start(start_indices[i]);
     slice_config->set_limit(limit_indices[i]);
     slice_config->set_stride(strides[i]);
+    if (i < start_exprs.size() && start_exprs[i] &&
+        start_exprs[i]->is_dynamic()) {
+      start_exprs[i].simplify().to_proto(slice_config->mutable_start_expr());
+    }
+    if (i < limit_exprs.size() && limit_exprs[i] &&
+        limit_exprs[i]->is_dynamic()) {
+      limit_exprs[i].simplify().to_proto(slice_config->mutable_limit_expr());
+    }
   }
   return AddInstruction(std::move(instr), HloOpcode::kSlice, {operand});
 }
@@ -1668,9 +1774,32 @@ XlaOp XlaBuilder::SliceInDim(XlaOp operand, int64_t start_index,
   });
 }
 
+XlaOp XlaBuilder::SliceInDim(XlaOp operand, int64_t start_index,
+                             int64_t limit_index, const DExpr& start_expr,
+                             const DExpr& limit_expr, int64_t stride,
+                             int64_t dimno) {
+  return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
+    TF_ASSIGN_OR_RETURN(const Shape* shape, GetShapePtr(operand));
+    std::vector<int64_t> starts(shape->dimensions().size(), 0);
+    std::vector<int64_t> limits(shape->dimensions().begin(),
+                                shape->dimensions().end());
+    std::vector<DExpr> start_exprs(shape->dimensions().size(), DExpr::Const(0));
+    std::vector<DExpr> limit_exprs(shape->expressions().begin(),
+                                   shape->expressions().end());
+    std::vector<int64_t> strides(shape->dimensions().size(), 1);
+    starts[dimno] = start_index;
+    limits[dimno] = limit_index;
+    start_exprs[dimno] = start_expr;
+    limit_exprs[dimno] = limit_expr;
+    strides[dimno] = stride;
+    return Slice(operand, starts, limits, start_exprs, limit_exprs, strides);
+  });
+}
+
 XlaOp XlaBuilder::DynamicSlice(XlaOp operand,
                                absl::Span<const XlaOp> start_indices,
-                               absl::Span<const int64_t> slice_sizes) {
+                               absl::Span<const int64_t> slice_sizes,
+                               absl::Span<const DExpr> slice_exprs) {
   return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
     std::vector<const Shape*> start_indices_shape_ptrs;
@@ -1679,9 +1808,9 @@ XlaOp XlaBuilder::DynamicSlice(XlaOp operand,
     absl::c_transform(start_indices_shapes,
                       std::back_inserter(start_indices_shape_ptrs),
                       [](const Shape& shape) { return &shape; });
-    TF_ASSIGN_OR_RETURN(Shape shape,
-                        ShapeInference::InferDynamicSliceShape(
-                            *operand_shape, start_indices_shapes, slice_sizes));
+    TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferDynamicSliceShape(
+                                         *operand_shape, start_indices_shapes,
+                                         slice_sizes, slice_exprs));
     return DynamicSliceInternal(shape, operand, start_indices, slice_sizes);
   });
 }
@@ -1698,6 +1827,7 @@ absl::StatusOr<XlaOp> XlaBuilder::DynamicSliceInternal(
 
   std::vector<XlaOp> operands = {operand};
   operands.insert(operands.end(), start_indices.begin(), start_indices.end());
+
   return AddInstruction(std::move(instr), HloOpcode::kDynamicSlice, operands);
 }
 
@@ -1794,9 +1924,22 @@ XlaOp XlaBuilder::Reshape(XlaOp operand, absl::Span<const int64_t> dimensions,
                           int64_t inferred_dimension) {
   return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
-    TF_ASSIGN_OR_RETURN(const Shape shape,
-                        ShapeInference::InferReshapeShape(
-                            *operand_shape, dimensions, inferred_dimension));
+    TF_ASSIGN_OR_RETURN(const Shape shape, ShapeInference::InferReshapeShape(
+                                               *operand_shape, dimensions,
+                                               inferred_dimension, {}));
+    return ReshapeInternal(shape, operand, inferred_dimension);
+  });
+}
+
+XlaOp XlaBuilder::Reshape(XlaOp operand, absl::Span<const int64_t> dimensions,
+                          absl::Span<const DExpr> expressions,
+                          int64_t inferred_dimension) {
+  return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
+    TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
+    TF_ASSIGN_OR_RETURN(
+        const Shape shape,
+        ShapeInference::InferReshapeShape(*operand_shape, dimensions,
+                                          inferred_dimension, expressions));
     return ReshapeInternal(shape, operand, inferred_dimension);
   });
 }
@@ -1811,19 +1954,19 @@ XlaOp XlaBuilder::Reshape(const Shape& shape, XlaOp operand,
 XlaOp XlaBuilder::DynamicReshape(XlaOp operand,
                                  absl::Span<const XlaOp> dim_sizes,
                                  absl::Span<const int64_t> new_size_bounds,
-                                 const std::vector<bool>& dims_are_dynamic) {
+                                 const std::vector<bool>& dims_are_dynamic,
+                                 absl::Span<const DExpr> expressions) {
   return ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
     std::vector<const Shape*> dim_size_shape_ptrs;
-    TF_ASSIGN_OR_RETURN(const auto& dim_size_shapes,
-                        GetOperandShapes(dim_sizes));
-
+    TF_ASSIGN_OR_RETURN(const auto& dim_size_shapes, GetOperandShapes(dim_sizes));
     absl::c_transform(dim_size_shapes, std::back_inserter(dim_size_shape_ptrs),
                       [](const Shape& shape) { return &shape; });
-    TF_ASSIGN_OR_RETURN(const Shape shape,
-                        ShapeInference::InferDynamicReshapeShape(
-                            *operand_shape, dim_size_shape_ptrs,
-                            new_size_bounds, dims_are_dynamic));
+    TF_ASSIGN_OR_RETURN(
+        const Shape shape,
+        ShapeInference::InferDynamicReshapeShape(
+            *operand_shape, dim_size_shape_ptrs, new_size_bounds,
+            dims_are_dynamic, expressions));
     TF_RETURN_IF_ERROR(first_error_);
     std::vector<XlaOp> operands;
     operands.reserve(1 + dim_sizes.size());
@@ -1865,17 +2008,20 @@ XlaOp XlaBuilder::Collapse(XlaOp operand,
     VLOG(3) << "dims to collapse: " << absl::StrJoin(dimensions, ",");
 
     std::vector<int64_t> new_sizes;
+    std::vector<DExpr> new_exprs;
     for (int i = 0; i < original_shape->dimensions().size(); ++i) {
       if (i <= dimensions.front() || i > dimensions.back()) {
         new_sizes.push_back(original_shape->dimensions(i));
+        new_exprs.push_back(original_shape->expressions(i));
       } else {
         new_sizes.back() *= original_shape->dimensions(i);
+        new_exprs.back() =
+            new_exprs.back() * original_shape->expressions(i);
       }
     }
 
     VLOG(3) << "new sizes: [" << absl::StrJoin(new_sizes, ",") << "]";
-
-    return Reshape(operand, new_sizes);
+    return Reshape(operand, new_sizes, new_exprs);
   });
 }
 
@@ -3912,12 +4058,14 @@ XlaOp XlaBuilder::AllToAllArray(
       return all_to_all;
     }
     DimensionVector sizes;
+    std::vector<DExpr> expressions;
     const bool is_unbounded = operand_shape->is_unbounded_dynamic();
     std::vector<XlaOp> dynamic_sizes;
     auto GetR1DimensionSizeOrConstant = [&](XlaOp operand,
                                             int64_t dimension) -> XlaOp {
       if (operand_shape->is_unbounded_dynamic_dimension(dimension)) {
-        return Reshape(GetDimensionSize(operand, dimension), {1});
+        return Reshape(GetDimensionSize(operand, dimension), {1},
+                       {DExpr::Const(1)});
       }
       return ConstantR1<int32_t>(
           this, {static_cast<int32_t>(operand_shape->dimensions(dimension))});
@@ -3927,15 +4075,18 @@ XlaOp XlaBuilder::AllToAllArray(
     for (int64_t i = 0; i < operand_shape->dimensions().size(); ++i) {
       if (i != split_dimension) {
         sizes.push_back(operand_shape->dimensions(i));
+        expressions.push_back(operand_shape->expressions(i));
         if (is_unbounded) {
           dynamic_sizes.push_back(GetR1DimensionSizeOrConstant(operand, i));
         }
         continue;
       }
       sizes.push_back(split_count);
+      expressions.push_back(DExpr::Const(split_count));
       sizes.push_back(operand_shape->is_unbounded_dynamic_dimension(i)
                           ? Shape::kUnboundedSize
                           : operand_shape->dimensions(i) / split_count);
+      expressions.push_back(operand_shape->expressions(i) / split_count);
 
       if (is_unbounded) {
         dynamic_sizes.push_back(r1_split_count);
@@ -3955,11 +4106,11 @@ XlaOp XlaBuilder::AllToAllArray(
       TF_ASSIGN_OR_RETURN(
           const Shape shape,
           ShapeUtil::MakeValidatedShape(all_to_all_shape.element_type(), sizes,
-                                        dynamic_dimensions));
+                                        dynamic_dimensions, expressions));
       all_to_all =
           MhloDynamicReshape(all_to_all, ConcatInDim(dynamic_sizes, 0), shape);
     } else {
-      all_to_all = Reshape(all_to_all, sizes);
+      all_to_all = Reshape(all_to_all, sizes, expressions);
     }
 
     std::vector<int64_t> permutation;
@@ -4489,6 +4640,14 @@ XlaOp XlaBuilder::GetDimensionSize(XlaOp operand, int64_t dimension) {
     TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand));
     TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferGetDimensionSizeShape(
                                          *operand_shape, dimension));
+    const DExpr& dim_expr = operand_shape->expressions(dimension);
+    if (dim_expr && dim_expr->is_dynamic()) {
+      DExpr simplified_expr = dim_expr.simplify();
+      XlaOp dim_bound =
+          ConstantR0<int32_t>(this, operand_shape->dimensions(dimension));
+      TF_RETURN_IF_ERROR(SetInstructionContents(dim_bound, {simplified_expr}));
+      return dim_bound;
+    }
     // Calling GetDimensionSize on a static dimension returns a constant
     // instruction.
     if (operand_shape->is_static_dimension(dimension)) {
@@ -4866,6 +5025,7 @@ absl::StatusOr<XlaOp> XlaBuilder::AddInstruction(
   TF_ASSIGN_OR_RETURN(Shape shape,
                       Shape::FromProto(instructions_.back().shape()));
   instruction_shapes_.push_back(std::make_unique<Shape>(std::move(shape)));
+  instruction_contents_.push_back({});
 
   XlaOp op(handle, this);
   return op;
@@ -4986,16 +5146,18 @@ XlaOp ConstantLiteral(XlaBuilder* builder, const LiteralSlice& literal) {
   return builder->ConstantLiteral(literal);
 }
 
-XlaOp Broadcast(const XlaOp operand,
-                absl::Span<const int64_t> broadcast_sizes) {
-  return operand.builder()->Broadcast(operand, broadcast_sizes);
+XlaOp Broadcast(const XlaOp operand, absl::Span<const int64_t> broadcast_sizes,
+                absl::Span<const DExpr> broadcast_exprs) {
+  return operand.builder()->Broadcast(operand, broadcast_sizes,
+                                      broadcast_exprs);
 }
 
 XlaOp BroadcastInDim(const XlaOp operand,
                      absl::Span<const int64_t> out_dim_size,
-                     absl::Span<const int64_t> broadcast_dimensions) {
+                     absl::Span<const int64_t> broadcast_dimensions,
+                     absl::Span<const DExpr> out_dim_exp) {
   return operand.builder()->BroadcastInDim(operand, out_dim_size,
-                                           broadcast_dimensions);
+                                           broadcast_dimensions, out_dim_exp);
 }
 
 XlaOp MhloDynamicReshape(const XlaOp operand, const XlaOp output_shape,
@@ -5030,21 +5192,38 @@ XlaOp Reshape(const XlaOp operand, absl::Span<const int64_t> dimensions) {
   return operand.builder()->Reshape(operand, dimensions);
 }
 
+XlaOp Reshape(const XlaOp operand, absl::Span<const int64_t> dimensions,
+              absl::Span<const DExpr> expressions) {
+  return operand.builder()->Reshape(operand, dimensions, expressions);
+}
+
 XlaOp Reshape(const Shape& shape, XlaOp operand) {
   return operand.builder()->Reshape(shape, operand);
 }
 
 XlaOp DynamicReshape(XlaOp operand, absl::Span<const XlaOp> dim_sizes,
                      absl::Span<const int64_t> new_size_bounds,
-                     const std::vector<bool>& dims_are_dynamic) {
+                     const std::vector<bool>& dims_are_dynamic,
+                     absl::Span<const DExpr> expressions) {
   return operand.builder()->DynamicReshape(operand, dim_sizes, new_size_bounds,
-                                           dims_are_dynamic);
+                                           dims_are_dynamic, expressions);
+}
+
+XlaOp Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
+            absl::Span<const int64_t> limit_indices,
+            absl::Span<const DExpr> start_exprs,
+            absl::Span<const DExpr> limit_exprs,
+            absl::Span<const int64_t> strides) {
+  return operand.builder()->Slice(operand, start_indices, limit_indices,
+                                  start_exprs, limit_exprs, strides);
 }
 
 XlaOp ReshapeWithInferredDimension(XlaOp operand,
                                    absl::Span<const int64_t> new_sizes,
+                                   absl::Span<const DExpr> new_exprs,
                                    int64_t inferred_dimension) {
-  return operand.builder()->Reshape(operand, new_sizes, inferred_dimension);
+  return operand.builder()->Reshape(operand, new_sizes, new_exprs,
+                                    inferred_dimension);
 }
 
 XlaOp Collapse(const XlaOp operand, absl::Span<const int64_t> dimensions) {
@@ -5064,9 +5243,18 @@ XlaOp SliceInDim(const XlaOp operand, int64_t start_index, int64_t limit_index,
                                        stride, dimno);
 }
 
+XlaOp SliceInDim(const XlaOp operand, int64_t start_index, int64_t limit_index,
+                 const DExpr& start_expr, const DExpr& limit_expr,
+                 int64_t stride, int64_t dimno) {
+  return operand.builder()->SliceInDim(operand, start_index, limit_index,
+                                       start_expr, limit_expr, stride, dimno);
+}
+
 XlaOp DynamicSlice(const XlaOp operand, absl::Span<const XlaOp> start_indices,
-                   absl::Span<const int64_t> slice_sizes) {
-  return operand.builder()->DynamicSlice(operand, start_indices, slice_sizes);
+                   absl::Span<const int64_t> slice_sizes,
+                   absl::Span<const DExpr> slice_exprs) {
+  return operand.builder()->DynamicSlice(operand, start_indices, slice_sizes,
+                                         slice_exprs);
 }
 
 XlaOp DynamicUpdateSlice(const XlaOp operand, const XlaOp update,
@@ -6085,6 +6273,14 @@ XlaOp GetDimensionSize(const XlaOp operand, int64_t dimension) {
   return operand.builder()->GetDimensionSize(operand, dimension);
 }
 
+XlaOp GetExpressionValue(XlaOp operand) {
+  XlaBuilder* builder = operand.builder();
+  return CustomCall(builder, "GetExpressionValue", {operand},
+                             ShapeUtil::MakeShape(S32, {}), "", false, {},
+                             nullptr, CustomCallSchedule::SCHEDULE_NONE,
+                             CustomCallApiVersion::API_VERSION_ORIGINAL);
+}
+
 XlaOp SetDimensionSize(const XlaOp operand, const XlaOp val,
                        int64_t dimension) {
   return operand.builder()->SetDimensionSize(operand, val, dimension);
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h
index 8f479090d86b19..b2c636c6eee530 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder.h
+++ b/third_party/xla/xla/hlo/builder/xla_builder.h
@@ -484,6 +484,13 @@ class XlaBuilder {
   absl::Status SetInstructionFrontendAttribute(XlaOp op, std::string attribute,
                                                std::string value);
 
+  // Associates symbolic contents metadata with a specific instruction.
+  absl::Status SetInstructionContents(XlaOp op, std::vector<DExpr> contents);
+
+  // Returns symbolic contents metadata attached to an instruction, if any.
+  absl::StatusOr<const std::vector<DExpr>*> GetInstructionContents(
+      XlaOp op) const;
+
   // Looks up the HloInstruction and sets the sharding. If the sharding already
   // existed, then its value is updated.
   //
@@ -519,10 +526,12 @@ class XlaBuilder {
 
   virtual XlaOp ConstantLiteral(const LiteralSlice& literal);
 
-  XlaOp Broadcast(XlaOp operand, absl::Span<const int64_t> broadcast_sizes);
+  XlaOp Broadcast(XlaOp operand, absl::Span<const int64_t> broadcast_sizes,
+                  absl::Span<const DExpr> broadcast_exprs = {});
 
   XlaOp BroadcastInDim(XlaOp operand, absl::Span<const int64_t> out_dim_size,
-                       absl::Span<const int64_t> broadcast_dimensions);
+                       absl::Span<const int64_t> broadcast_dimensions,
+                       absl::Span<const DExpr> out_dim_exp = {});
 
   // This is an experimental API for creating the mhlo.dynamic_broadcast_in_dim
   // op from the XlaBuilder. This is only intended for export to MHLO or
@@ -545,12 +554,17 @@ class XlaBuilder {
   XlaOp Reshape(XlaOp operand, absl::Span<const int64_t> dimensions,
                 int64_t inferred_dimension = -1);
 
+  XlaOp Reshape(XlaOp operand, absl::Span<const int64_t> dimensions,
+                absl::Span<const DExpr> expressions,
+                int64_t inferred_dimension = -1);
+
   XlaOp Reshape(const Shape& shape, XlaOp operand,
                 int64_t inferred_dimension = -1);
 
   XlaOp DynamicReshape(XlaOp operand, absl::Span<const XlaOp> dim_sizes,
                        absl::Span<const int64_t> new_size_bounds,
-                       const std::vector<bool>& dims_are_dynamic);
+                       const std::vector<bool>& dims_are_dynamic,
+                       absl::Span<const DExpr> expressions = {});
 
   XlaOp MhloDynamicReshape(XlaOp operand, XlaOp output_shape,
                            const Shape& shape);
@@ -560,16 +574,36 @@ class XlaBuilder {
   XlaOp Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
               absl::Span<const int64_t> limit_indices,
               absl::Span<const int64_t> strides);
+
+  XlaOp Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
+              absl::Span<const int64_t> limit_indices,
+              absl::Span<const DExpr> start_exprs,
+              absl::Span<const DExpr> limit_exprs,
+              absl::Span<const int64_t> strides);
+
   virtual absl::StatusOr<XlaOp> SliceInternal(
       const Shape& shape, XlaOp operand,
       absl::Span<const int64_t> start_indices,
       absl::Span<const int64_t> limit_indices,
       absl::Span<const int64_t> strides);
+  virtual absl::StatusOr<XlaOp> SliceInternal(
+      const Shape& shape, XlaOp operand,
+      absl::Span<const int64_t> start_indices,
+      absl::Span<const int64_t> limit_indices,
+      absl::Span<const DExpr> start_exprs,
+      absl::Span<const DExpr> limit_exprs,
+      absl::Span<const int64_t> strides);
   virtual XlaOp SliceInDim(XlaOp operand, int64_t start_index,
                            int64_t limit_index, int64_t stride, int64_t dimno);
 
+  virtual XlaOp SliceInDim(XlaOp operand, int64_t start_index,
+                           int64_t limit_index, const DExpr& start_expr,
+                           const DExpr& limit_expr, int64_t stride,
+                           int64_t dimno);
+
   XlaOp DynamicSlice(XlaOp operand, absl::Span<const XlaOp> start_indices,
-                     absl::Span<const int64_t> slice_sizes);
+                     absl::Span<const int64_t> slice_sizes,
+                     absl::Span<const DExpr> slice_exprs = {});
   virtual absl::StatusOr<XlaOp> DynamicSliceInternal(
       const Shape& shape, XlaOp operand, absl::Span<const XlaOp> start_indices,
       absl::Span<const int64_t> slice_sizes);
@@ -1181,6 +1215,7 @@ class XlaBuilder {
   // A cache for the HloInstructionProto shapes, to avoid recreating Shape
   // objects from protos and to support the GetShapePtr() API.
   std::vector<std::unique_ptr<Shape>> instruction_shapes_;
+  std::vector<std::vector<DExpr>> instruction_contents_;
 
   // Dynamic parameter configuration of this computation.
   DynamicParameterBinding dynamic_parameter_binding_;
@@ -1244,11 +1279,13 @@ class XlaBuilder {
                                const LiteralSlice& literal);
 
   friend XlaOp Broadcast(XlaOp operand,
-                         absl::Span<const int64_t> broadcast_sizes);
+                         absl::Span<const int64_t> broadcast_sizes,
+                         absl::Span<const DExpr> broadcast_expressions);
 
   friend XlaOp BroadcastInDim(XlaOp operand,
                               absl::Span<const int64_t> out_dim_size,
-                              absl::Span<const int64_t> broadcast_dimensions);
+                              absl::Span<const int64_t> broadcast_dimensions,
+                              absl::Span<const DExpr> out_dim_exp);
 
   friend XlaOp MhloDynamicBroadcastInDim(
       XlaOp operand, XlaOp output_dimensions,
@@ -1265,18 +1302,22 @@ class XlaBuilder {
 
   friend XlaOp Reshape(XlaOp operand, absl::Span<const int64_t> dimensions);
 
+  friend XlaOp Reshape(XlaOp operand, absl::Span<const int64_t> dimensions,
+                       absl::Span<const DExpr> expressions);
+
   friend XlaOp Reshape(const Shape& shape, XlaOp operand);
 
   friend XlaOp DynamicReshape(XlaOp operand, absl::Span<const XlaOp> dim_sizes,
                               absl::Span<const int64_t> new_size_bounds,
-                              const std::vector<bool>& dims_are_dynamic);
+                              const std::vector<bool>& dims_are_dynamic,
+                              absl::Span<const DExpr> expressions);
 
   friend XlaOp MhloDynamicReshape(XlaOp operand, XlaOp output_shape,
                                   const Shape& shape);
 
-  friend XlaOp ReshapeWithInferredDimension(XlaOp operand,
-                                            absl::Span<const int64_t> new_sizes,
-                                            int64_t inferred_dimension);
+  friend XlaOp ReshapeWithInferredDimension(
+      XlaOp operand, absl::Span<const int64_t> new_sizes,
+      absl::Span<const DExpr> new_exprs, int64_t inferred_dimension);
 
   friend XlaOp Collapse(XlaOp operand, absl::Span<const int64_t> dimensions);
 
@@ -1284,12 +1325,24 @@ class XlaBuilder {
                      absl::Span<const int64_t> limit_indices,
                      absl::Span<const int64_t> strides);
 
+  friend XlaOp Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
+                     absl::Span<const int64_t> limit_indices,
+                     absl::Span<const DExpr> start_exprs,
+                     absl::Span<const DExpr> limit_exprs,
+                     absl::Span<const int64_t> strides);
+
   friend XlaOp SliceInDim(XlaOp operand, int64_t start_index,
                           int64_t limit_index, int64_t stride, int64_t dimno);
 
+  friend XlaOp SliceInDim(XlaOp operand, int64_t start_index,
+                          int64_t limit_index, const DExpr& start_expr,
+                          const DExpr& limit_expr, int64_t stride,
+                          int64_t dimno);
+
   friend XlaOp DynamicSlice(XlaOp operand,
                             absl::Span<const XlaOp> start_indices,
-                            absl::Span<const int64_t> slice_sizes);
+                            absl::Span<const int64_t> slice_sizes,
+                            absl::Span<const DExpr> slice_exprs);
 
   friend XlaOp DynamicUpdateSlice(XlaOp operand, XlaOp update,
                                   absl::Span<const XlaOp> start_indices);
@@ -1975,7 +2028,8 @@ XlaOp ConstantR1(XlaBuilder* builder, int64_t length, NativeT value);
 // The new dimensions index into copies of the operand, i.e.
 //
 //   output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]
-XlaOp Broadcast(XlaOp operand, absl::Span<const int64_t> broadcast_sizes);
+XlaOp Broadcast(XlaOp operand, absl::Span<const int64_t> broadcast_sizes,
+                absl::Span<const DExpr> broadcast_exprs = {});
 
 // This op broadcasts the `operand` to an output with the given `shape`.
 // `broadcast_dimensions` are the dimensions to be broadcasting into, i.e., the
@@ -1993,7 +2047,8 @@ XlaOp Broadcast(XlaOp operand, absl::Span<const int64_t> broadcast_sizes);
 //   {{1 , 1},
 //    {2 , 2}}
 XlaOp BroadcastInDim(XlaOp operand, absl::Span<const int64_t> out_dim_size,
-                     absl::Span<const int64_t> broadcast_dimensions);
+                     absl::Span<const int64_t> broadcast_dimensions,
+                     absl::Span<const DExpr> out_dim_exp = {});
 
 // This is an experimental API for creating the mhlo.dynamic_broadcast_in_dim
 // op from the XlaBuilder. This is only intended for export to MHLO or
@@ -2038,7 +2093,8 @@ XlaOp PadInDim(XlaOp operand, XlaOp padding_value, int64_t dimno,
 // dimension dimension if dims_are_dynamic[i] is true.
 XlaOp DynamicReshape(XlaOp operand, absl::Span<const XlaOp> dim_sizes,
                      absl::Span<const int64_t> new_size_bounds,
-                     const std::vector<bool>& dims_are_dynamic);
+                     const std::vector<bool>& dims_are_dynamic,
+                     absl::Span<const DExpr> expressions = {});
 
 // This is an experimental API for creating the mhlo.dynamic_reshape op from the
 // XlaBuilder. This is only intended for export to MHLO or StableHLO, and cannot
@@ -2050,6 +2106,9 @@ XlaOp MhloDynamicReshape(XlaOp operand, XlaOp output_shape, const Shape& shape);
 // dimension sizes. Conceptually, this is a limited form of "shape casting".
 XlaOp Reshape(XlaOp operand, absl::Span<const int64_t> dimensions);
 
+XlaOp Reshape(XlaOp operand, absl::Span<const int64_t> dimensions,
+              absl::Span<const DExpr> expressions);
+
 // Enqueues a Reshape op that uses an explicit target shape.
 XlaOp Reshape(const Shape& shape, XlaOp operand);
 
@@ -2059,6 +2118,7 @@ XlaOp Reshape(const Shape& shape, XlaOp operand);
 // is a dynamic dimension in the output, it must be the inferred dimension.
 XlaOp ReshapeWithInferredDimension(XlaOp operand,
                                    absl::Span<const int64_t> new_sizes,
+                                   absl::Span<const DExpr> new_exprs,
                                    int64_t inferred_dimension);
 
 // Wrapper for Reshape.
@@ -2096,6 +2156,12 @@ XlaOp Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
             absl::Span<const int64_t> limit_indices,
             absl::Span<const int64_t> strides);
 
+XlaOp Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
+            absl::Span<const int64_t> limit_indices,
+            absl::Span<const DExpr> start_exprs,
+            absl::Span<const DExpr> limit_exprs,
+            absl::Span<const int64_t> strides);
+
 // Enqueues a slice operation in a given dimension, taking all other
 // dimensions as they are; e.g. if dimno is 1 from start_index 2 to
 // limit_index 4 by 1, and the shape is f32[7,8,9], this call is short-hand
@@ -2105,6 +2171,10 @@ XlaOp Slice(XlaOp operand, absl::Span<const int64_t> start_indices,
 XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index,
                  int64_t stride, int64_t dimno);
 
+XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index,
+                 const DExpr& start_expr, const DExpr& limit_expr,
+                 int64_t stride, int64_t dimno);
+
 // Enqueues a slice operation onto the computation that slices the 'operand'
 // from dynamic start indices which are passed in 'start_indices'.
 // The size of the slice in each dimension is passed in 'slice_sizes',
@@ -2116,7 +2186,8 @@ XlaOp SliceInDim(XlaOp operand, int64_t start_index, int64_t limit_index,
 // Slice index calculations are computed modulo input dimension sizes to
 // prevent dynamic start indices from generating out-of-bound array accesses.
 XlaOp DynamicSlice(XlaOp operand, absl::Span<const XlaOp> start_indices,
-                   absl::Span<const int64_t> slice_sizes);
+                   absl::Span<const int64_t> slice_sizes,
+                   absl::Span<const DExpr> slice_exprs = {});
 
 // Enqueues a dynamic update slice operation onto the computation, which
 // updates a slice of 'operand' with 'update' at dynamic 'start_indices'.
@@ -3055,6 +3126,8 @@ XlaOp BatchNormGrad(XlaOp operand, XlaOp scale, XlaOp batch_mean,
 // array shaped.
 XlaOp GetDimensionSize(XlaOp operand, int64_t dimension);
 
+XlaOp GetExpressionValue(XlaOp operand);
+
 // Sets the size of the given dimension of the operand. The operand must be
 // array shaped.  The result will have the same shape as the operand, but the
 // given dimension will be dynamic (if not already).
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
index 0dc7d619afc370..86b7bb0facc892 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
@@ -108,6 +108,64 @@ absl::Status EraseElementFromVector(PtrVec<T>* container, T value) {
   container->erase(it);
   return absl::OkStatus();
 }
+
+DynExpr* DynExprFromProtoForPrint(const ExpressionProto& proto) {
+  switch (proto.node_type_case()) {
+    case ExpressionProto::kConstantValue:
+      return DynExpr::_(proto.constant_value());
+    case ExpressionProto::kVariableId:
+      return DynExpr::V(proto.variable_id());
+    case ExpressionProto::kAddNode: {
+      const auto& add = proto.add_node();
+      return new Add(DynExprFromProtoForPrint(add.lhs()),
+                     DynExprFromProtoForPrint(add.rhs()));
+    }
+    case ExpressionProto::kSubNode: {
+      const auto& sub = proto.sub_node();
+      return new Sub(DynExprFromProtoForPrint(sub.lhs()),
+                     DynExprFromProtoForPrint(sub.rhs()));
+    }
+    case ExpressionProto::kMulNode: {
+      const auto& mul = proto.mul_node();
+      return new Mul(DynExprFromProtoForPrint(mul.lhs()),
+                     DynExprFromProtoForPrint(mul.rhs()));
+    }
+    case ExpressionProto::kDivNode: {
+      const auto& div = proto.div_node();
+      return new Div(DynExprFromProtoForPrint(div.lhs()),
+                     DynExprFromProtoForPrint(div.rhs()));
+    }
+    case ExpressionProto::kMaxNode: {
+      const auto& max = proto.max_node();
+      return new MaxExpr(DynExprFromProtoForPrint(max.lhs()),
+                         DynExprFromProtoForPrint(max.rhs()));
+    }
+    case ExpressionProto::kGtNode: {
+      const auto& gt = proto.gt_node();
+      return new GtExpr(DynExprFromProtoForPrint(gt.lhs()),
+                        DynExprFromProtoForPrint(gt.rhs()));
+    }
+    case ExpressionProto::kSelectNode: {
+      const auto& select = proto.select_node();
+      return new SelectExpr(DynExprFromProtoForPrint(select.pred()),
+                            DynExprFromProtoForPrint(select.on_true()),
+                            DynExprFromProtoForPrint(select.on_false()));
+    }
+    case ExpressionProto::NODE_TYPE_NOT_SET:
+    default:
+      return nullptr;
+  }
+}
+
+std::string ContentsExprToString(const ExpressionProto& proto) {
+  std::unique_ptr<DynExpr> expr(DynExprFromProtoForPrint(proto));
+  if (expr == nullptr) {
+    return "_";
+  }
+  StringPrinter printer;
+  expr->print(&printer);
+  return std::move(printer).ToString();
+}
 }  // namespace
 
 HloInstruction::Users::~Users() = default;
@@ -636,14 +694,32 @@ absl::StatusOr<std::unique_ptr<HloInstruction>> HloInstruction::CreateFromProto(
       break;
     case HloOpcode::kSlice: {
       std::vector<int64_t> slice_starts, slice_limits, slice_strides;
+      std::vector<DExpr> slice_start_exprs, slice_limit_exprs;
+      bool has_symbolic_slice_bounds = false;
       for (const HloInstructionProto::SliceDimensions& slice_dimensions :
            proto.slice_dimensions()) {
         slice_starts.push_back(slice_dimensions.start());
         slice_limits.push_back(slice_dimensions.limit());
         slice_strides.push_back(slice_dimensions.stride());
+        if (slice_dimensions.has_start_expr() ||
+            slice_dimensions.has_limit_expr()) {
+          has_symbolic_slice_bounds = true;
+        }
+        slice_start_exprs.push_back(
+            slice_dimensions.has_start_expr()
+                ? DExprFromProto(slice_dimensions.start_expr())
+                : DExpr::Unknown(kMissingExpressionSentinel));
+        slice_limit_exprs.push_back(
+            slice_dimensions.has_limit_expr()
+                ? DExprFromProto(slice_dimensions.limit_expr())
+                : DExpr::Unknown(kMissingExpressionSentinel));
       }
-      instruction = CreateSlice(shape, operands(0), slice_starts, slice_limits,
-                                slice_strides);
+      instruction = has_symbolic_slice_bounds
+                        ? CreateSlice(shape, operands(0), slice_starts,
+                                      slice_limits, slice_strides,
+                                      slice_start_exprs, slice_limit_exprs)
+                        : CreateSlice(shape, operands(0), slice_starts,
+                                      slice_limits, slice_strides);
       break;
     }
     case HloOpcode::kConstant: {
@@ -1373,6 +1449,14 @@ absl::StatusOr<std::unique_ptr<HloInstruction>> HloInstruction::CreateFromProto(
   if (proto.has_frontend_attributes()) {
     instruction->set_frontend_attributes(proto.frontend_attributes());
   }
+  if (proto.contents_size() > 0) {
+    std::vector<ExpressionProto> contents;
+    contents.reserve(proto.contents_size());
+    for (const auto& content : proto.contents()) {
+      contents.push_back(content);
+    }
+    instruction->set_contents(std::move(contents));
+  }
 
   if (proto.has_statistics_viz()) {
     instruction->set_statistics_viz(proto.statistics_viz());
@@ -2046,6 +2130,18 @@ HloInstruction::CreateAddDependency(HloInstruction* data_operand,
                                                limit_indices, strides);
 }
 
+/* static */ std::unique_ptr<HloInstruction> HloInstruction::CreateSlice(
+    const Shape& shape, HloInstruction* operand,
+    absl::Span<const int64_t> start_indices,
+    absl::Span<const int64_t> limit_indices,
+    absl::Span<const int64_t> strides,
+    absl::Span<const DExpr> start_exprs,
+    absl::Span<const DExpr> limit_exprs) {
+  return std::make_unique<HloSliceInstruction>(shape, operand, start_indices,
+                                               limit_indices, strides,
+                                               start_exprs, limit_exprs);
+}
+
 /* static */ std::unique_ptr<HloInstruction> HloInstruction::CreateDynamicSlice(
     const Shape& shape, HloInstruction* operand,
     absl::Span<HloInstruction* const> start_indices,
@@ -2812,6 +2908,7 @@ std::unique_ptr<HloInstruction> HloInstruction::CloneWithNewOperands(
   SetupDerivedInstruction(clone.get());
   clone->backend_config_ = BackendConfigWrapper(backend_config_);
   clone->set_frontend_attributes(frontend_attributes());
+  clone->set_contents(contents());
   // The new instruction's name will be uniquified when it's added to a
   // computation.
   clone->SetAndSanitizeName(name());
@@ -4235,6 +4332,18 @@ void HloInstruction::PrintExtraAttributes(
                 FrontendAttributesToString(frontend_attributes()));
     });
   }
+  if (has_contents()) {
+    printer.Next([this](Printer* printer) {
+      printer->Append("contents=[");
+      for (int64_t i = 0; i < contents().size(); ++i) {
+        if (i > 0) {
+          printer->Append(", ");
+        }
+        printer->Append(ContentsExprToString(contents()[i]));
+      }
+      printer->Append("]");
+    });
+  }
 
   if (opcode() != HloOpcode::kCall) {
     CHECK(!is_composite())
@@ -4356,6 +4465,9 @@ HloInstructionProto HloInstruction::ToProto() const {
   }
 
   *proto.mutable_frontend_attributes() = frontend_attributes();
+  for (const auto& content : contents()) {
+    *proto.add_contents() = content;
+  }
   proto.set_is_composite(is_composite());
 
   *proto.mutable_statistics_viz() = statistics_viz();
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h
index 59f542b6ec3b5b..8feebcc1647a3c 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction.h
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h
@@ -865,6 +865,13 @@ class HloInstruction {
       absl::Span<const int64_t> start_indices,
       absl::Span<const int64_t> limit_indices,
       absl::Span<const int64_t> strides);
+  static std::unique_ptr<HloInstruction> CreateSlice(
+      const Shape& shape, HloInstruction* operand,
+      absl::Span<const int64_t> start_indices,
+      absl::Span<const int64_t> limit_indices,
+      absl::Span<const int64_t> strides,
+      absl::Span<const DExpr> start_exprs,
+      absl::Span<const DExpr> limit_exprs);
 
   // Creates a slice instruction, where the first operand is sliced by
   // start indices specified in the second operand, and by size specified in
@@ -1883,6 +1890,19 @@ class HloInstruction {
     return rare()->frontend_attributes;
   }
 
+  void set_contents(std::vector<ExpressionProto> contents) {
+    if (!has_rare() && contents.empty()) {
+      return;
+    }
+    mutable_rare()->contents = std::move(contents);
+  }
+
+  const std::vector<ExpressionProto>& contents() const {
+    return rare()->contents;
+  }
+
+  bool has_contents() const { return has_rare() && !rare()->contents.empty(); }
+
   std::optional<std::string> get_frontend_attribute(
       absl::string_view key) const {
     auto it = rare()->frontend_attributes.map().find(key);
@@ -2506,6 +2526,9 @@ class HloInstruction {
     //    z' = const(20), frontend_attributes={?}
     FrontendAttributes frontend_attributes;
 
+    // Structured symbolic contents attached to this instruction.
+    std::vector<ExpressionProto> contents;
+
     // Used by kCall to determine if the Call instruction is a composite.
     bool is_composite;
 
diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc
index 0557a52c67fd1d..e4853f8c56a4fb 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc
@@ -74,6 +74,15 @@ namespace {
 using absl::CEscape;
 using absl::StrCat;
 
+std::string DExprToString(const DExpr& expr) {
+  if (!expr) {
+    return "_";
+  }
+  StringPrinter printer;
+  expr.simplify()->print(&printer);
+  return std::move(printer).ToString();
+}
+
 bool IsInstructionElementwiseOnOperand(const HloInstruction* instruction,
                                        const HloInstruction* operand) {
   const auto operand_indices = instruction->OperandIndices(operand);
@@ -1651,10 +1660,20 @@ HloSliceInstruction::HloSliceInstruction(
     const Shape& shape, HloInstruction* operand,
     absl::Span<const int64_t> start_indices,
     absl::Span<const int64_t> limit_indices, absl::Span<const int64_t> strides)
+    : HloSliceInstruction(shape, operand, start_indices, limit_indices, strides,
+                          /*start_exprs=*/{}, /*limit_exprs=*/{}) {}
+
+HloSliceInstruction::HloSliceInstruction(
+    const Shape& shape, HloInstruction* operand,
+    absl::Span<const int64_t> start_indices,
+    absl::Span<const int64_t> limit_indices, absl::Span<const int64_t> strides,
+    absl::Span<const DExpr> start_exprs, absl::Span<const DExpr> limit_exprs)
     : HloInstruction(HloOpcode::kSlice, shape),
       slice_starts_(start_indices.begin(), start_indices.end()),
       slice_limits_(limit_indices.begin(), limit_indices.end()),
-      slice_strides_(strides.begin(), strides.end()) {
+      slice_strides_(strides.begin(), strides.end()),
+      slice_start_exprs_(start_exprs.begin(), start_exprs.end()),
+      slice_limit_exprs_(limit_exprs.begin(), limit_exprs.end()) {
   AppendOperand(operand);
   // For backward compatibility with old serialized computations: if there are
   // no strides, assume all strides are 1.
@@ -1671,6 +1690,14 @@ HloInstructionProto HloSliceInstruction::ToProto() const {
     slice_dimension->set_start(slice_starts_[i]);
     slice_dimension->set_limit(slice_limits_[i]);
     slice_dimension->set_stride(slice_strides_[i]);
+    const DExpr& start_expr = slice_start_exprs(i);
+    const DExpr& limit_expr = slice_limit_exprs(i);
+    if (start_expr) {
+      start_expr.to_proto(slice_dimension->mutable_start_expr());
+    }
+    if (limit_expr) {
+      limit_expr.to_proto(slice_dimension->mutable_limit_expr());
+    }
   }
   return proto;
 }
@@ -1688,6 +1715,12 @@ void HloSliceInstruction::PrintExtraAttributesImpl(
                  if (!omit_stride) {
                    AppendCat(printer, ":", slice_strides_[i]);
                  }
+                 const DExpr& start_expr = slice_start_exprs(i);
+                 const DExpr& limit_expr = slice_limit_exprs(i);
+                 if (start_expr || limit_expr) {
+                   AppendCat(printer, " start_expr=", DExprToString(start_expr),
+                             " limit_expr=", DExprToString(limit_expr));
+                 }
                  printer->Append("]");
                });
     printer->Append("}");
@@ -1701,7 +1734,9 @@ bool HloSliceInstruction::IdenticalSlowPath(
   const auto& other_slice = static_cast<const HloSliceInstruction&>(other);
   return slice_starts_ == other_slice.slice_starts_ &&
          slice_limits_ == other_slice.slice_limits_ &&
-         slice_strides_ == other_slice.slice_strides_;
+         slice_strides_ == other_slice.slice_strides_ &&
+         slice_start_exprs_ == other_slice.slice_start_exprs_ &&
+         slice_limit_exprs_ == other_slice.slice_limit_exprs_;
 }
 
 std::unique_ptr<HloInstruction> HloSliceInstruction::CloneWithNewOperandsImpl(
@@ -1709,7 +1744,8 @@ std::unique_ptr<HloInstruction> HloSliceInstruction::CloneWithNewOperandsImpl(
     HloCloneContext* context) const {
   CHECK_EQ(new_operands.size(), 1);
   return std::make_unique<HloSliceInstruction>(
-      shape, new_operands[0], slice_starts_, slice_limits_, slice_strides_);
+      shape, new_operands[0], slice_starts_, slice_limits_, slice_strides_,
+      slice_start_exprs_, slice_limit_exprs_);
 }
 
 HloConstantInstruction::HloConstantInstruction(Literal literal)
diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h
index 22dd6ebe0e7b50..e90399a66a4ba9 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instructions.h
+++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h
@@ -1250,6 +1250,12 @@ class HloSliceInstruction : public HloInstruction {
                                absl::Span<const int64_t> start_indices,
                                absl::Span<const int64_t> limit_indices,
                                absl::Span<const int64_t> strides);
+  explicit HloSliceInstruction(const Shape& shape, HloInstruction* operand,
+                               absl::Span<const int64_t> start_indices,
+                               absl::Span<const int64_t> limit_indices,
+                               absl::Span<const int64_t> strides,
+                               absl::Span<const DExpr> start_exprs,
+                               absl::Span<const DExpr> limit_exprs);
 
   HloInstructionProto ToProto() const override;
 
@@ -1275,6 +1281,32 @@ class HloSliceInstruction : public HloInstruction {
   const std::vector<int64_t>& slice_strides() const { return slice_strides_; }
   std::vector<int64_t>* mutable_slice_strides() { return &slice_strides_; }
 
+  const DExpr& slice_start_exprs(int64_t dimension) const {
+    static const DExpr kMissingSliceExpr =
+        DExpr::Unknown(kMissingExpressionSentinel);
+    if (dimension < 0 ||
+        dimension >= static_cast<int64_t>(slice_start_exprs_.size())) {
+      return kMissingSliceExpr;
+    }
+    return slice_start_exprs_[dimension];
+  }
+  const std::vector<DExpr>& slice_start_exprs() const {
+    return slice_start_exprs_;
+  }
+
+  const DExpr& slice_limit_exprs(int64_t dimension) const {
+    static const DExpr kMissingSliceExpr =
+        DExpr::Unknown(kMissingExpressionSentinel);
+    if (dimension < 0 ||
+        dimension >= static_cast<int64_t>(slice_limit_exprs_.size())) {
+      return kMissingSliceExpr;
+    }
+    return slice_limit_exprs_[dimension];
+  }
+  const std::vector<DExpr>& slice_limit_exprs() const {
+    return slice_limit_exprs_;
+  }
+
   static bool ClassOf(const HloInstruction* hlo) {
     return hlo->opcode() == HloOpcode::kSlice;
   }
@@ -1295,6 +1327,8 @@ class HloSliceInstruction : public HloInstruction {
   std::vector<int64_t> slice_starts_;
   std::vector<int64_t> slice_limits_;
   std::vector<int64_t> slice_strides_;
+  std::vector<DExpr> slice_start_exprs_;
+  std::vector<DExpr> slice_limit_exprs_;
 };
 
 class HloConstantInstruction : public HloInstruction {
diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc
index 07cf9a8c4bce2c..11eeefa948fa1b 100644
--- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc
+++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc
@@ -250,6 +250,7 @@ absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
       }
       TF_RETURN_IF_ERROR(status);
     }
+
     if (!pass->IsPassPipeline()) {
       compilation_stats_->EndPass(pass_name);
     }
diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc b/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc
index 61b261248e0e2c..aa781b082c8823 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_combiner.cc
@@ -124,9 +124,9 @@ absl::Status CombineAllGathers(absl::Span<HloInstruction* const> to_combine,
                 (*perm)[ag->all_gather_dimension()]);
 
       // Bitcast operand and update output shape.
+      auto sh = ShapeUtil::PermuteDimensions(*perm, operand_shape);
       operands.back() =
-          computation.AddInstruction(HloInstruction::CreateBitcast(
-              ShapeUtil::PermuteDimensions(*perm, operand_shape), operand));
+          computation.AddInstruction(HloInstruction::CreateBitcast(sh, operand));
       output_shapes.back() = ShapeUtil::PermuteDimensions(*perm, hlo->shape());
     }
   }
diff --git a/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc b/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc
index b3c2ffe79ec00c..e038b7338a1bc2 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/collective_quantizer.cc
@@ -121,6 +121,7 @@ HloInstruction* ApplyUnaries(HloInstruction* instr,
     instr = instr->AddInstruction(unary->CloneWithNewOperands(
         ShapeUtil::MakeShapeWithDenseLayout(
             instr->shape().element_type(), unary->shape().dimensions(),
+            unary->shape().expressions(),
             unary->shape().layout().minor_to_major()),
         {instr}));
   }
diff --git a/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc
index ddb505801d92ca..ecdc47c0877542 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/bitcast_dtypes_expander.cc
@@ -80,10 +80,13 @@ absl::StatusOr<HloInstruction*> BitcastDtypesExpander::ExpandInstruction(
       broadcasted_input_shape.push_back(input_bit_width / output_bit_width);
       reshaped_input_shape.push_back(1);
       int64_t output_bit_width_mask = (int64_t{1} << output_bit_width) - 1;
-
-      TF_ASSIGN_OR_RETURN(input,
-                          BroadcastTo(Reshape(input, reshaped_input_shape),
-                                      broadcasted_input_shape));
+      std::vector<DExpr> reshaped_input_exprs(
+          from_shape.expressions().begin(), from_shape.expressions().end());
+      reshaped_input_exprs.push_back(DExpr::Const(1));
+      TF_ASSIGN_OR_RETURN(
+          input, BroadcastTo(
+                     Reshape(input, reshaped_input_shape, reshaped_input_exprs),
+                     broadcasted_input_shape));
       input = BitcastConvertType(input, input_logical_type);
       TF_ASSIGN_OR_RETURN(Shape input_shape, b.GetShape(input));
       XlaOp iota = Iota(&b, input_shape, input_shape.dimensions().size() - 1);
diff --git a/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc
index b5df50e1956c90..7b11a211ee3e02 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/cholesky_expander.cc
@@ -218,9 +218,9 @@ XlaOp CholeskyExpander::BuildCholesky(XlaOp a, int64_t block_size,
         l = UpdateSliceInMinorDims(l, update, {i + k, i});
       }
     }
-    return Select(
-        BroadcastInDim(seen_error, a_shape.dimensions(), error_dim_indices),
-        FullLike(l, std::numeric_limits<float>::quiet_NaN()), l);
+    return Select(BroadcastInDim(seen_error, a_shape.dimensions(),
+                                 error_dim_indices, a_shape.expressions()),
+                  FullLike(l, std::numeric_limits<float>::quiet_NaN()), l);
   });
 }
 
diff --git a/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc b/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc
index a3787d88ebbd93..e1aa51a390e2a0 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/dot_decomposer.cc
@@ -80,26 +80,42 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) {
   lhs_non_contracting_dims.reserve(num_lhs_non_contracting_dims);
   int64_t lhs_contracting_size = 1;
   bool lhs_contracting_dynamic = false;
+  int64_t lhs_contracting_multiplier_accu = 1;
+  DExpr lhs_contracting_expression = DExpr::Const(1);
   int64_t lhs_non_contracting_size = 1;
   bool lhs_non_contracting_dynamic = false;
+  int64_t lhs_non_contracting_multiplier_accu = 1;
+  DExpr lhs_non_contracting_expression = DExpr::Const(1);
   std::vector<int64_t> batch_dim_sizes;
   batch_dim_sizes.reserve(num_batch_dims);
   std::vector<bool> batch_dynamic_dims;
   batch_dynamic_dims.reserve(num_batch_dims);
+  std::vector<DExpr> batch_expressions;
+  batch_expressions.reserve(num_batch_dims);
+
+  bool lhs_contracting_is_static = true;
+  bool lhs_non_contracting_is_static = true;
+
   for (int64_t i = 0; i < lhs_rank; ++i) {
     if (absl::c_linear_search(original_dnums.lhs_contracting_dimensions(), i)) {
       lhs_contracting_size *= lhs_shape.dimensions(i);
       lhs_contracting_dynamic |= lhs_shape.is_dynamic_dimension(i);
+      lhs_contracting_expression =
+          lhs_contracting_expression * lhs_shape.expressions(i);
     } else if (absl::c_linear_search(original_dnums.lhs_batch_dimensions(),
                                      i)) {
       batch_dim_sizes.push_back(lhs_shape.dimensions(i));
       batch_dynamic_dims.push_back(lhs_shape.is_dynamic_dimension(i));
+      batch_expressions.push_back(lhs_shape.expressions(i));
     } else {
       lhs_non_contracting_dims.push_back(i);
       lhs_non_contracting_size *= lhs_shape.dimensions(i);
       lhs_non_contracting_dynamic |= lhs_shape.is_dynamic_dimension(i);
+      lhs_non_contracting_expression =
+          lhs_non_contracting_expression * lhs_shape.expressions(i);
     }
   }
+
   // The canonical form of the lhs is
   // [BatchDims, NonContractingDimsProduct, ContractingsDimsProduct]
   // If NonContractingDimsProduct is 1, it is omitted.
@@ -123,18 +139,21 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) {
 
   std::vector<int64_t> lhs_reshape_dims = batch_dim_sizes;
   std::vector<bool> lhs_reshape_dynamic_dims = batch_dynamic_dims;
+  std::vector<DExpr> lhs_reshape_expressions = batch_expressions;
   if (lhs_non_contracting_size > 1) {
     lhs_reshape_dims.push_back(lhs_non_contracting_size);
     lhs_reshape_dynamic_dims.push_back(lhs_non_contracting_dynamic);
+    lhs_reshape_expressions.push_back(lhs_non_contracting_expression);
   }
   lhs_reshape_dims.push_back(lhs_contracting_size);
   lhs_reshape_dynamic_dims.push_back(lhs_contracting_dynamic);
+  lhs_reshape_expressions.push_back(lhs_contracting_expression);
   // Reshape the contracting and non-contracting dimensions together.
+  auto sh_lhs = ShapeUtil::MakeShape(lhs_shape.element_type(), lhs_reshape_dims,
+                                     lhs_reshape_dynamic_dims,
+                                     lhs_reshape_expressions);
   HloInstruction* reshaped_lhs = computation->AddInstruction(
-      HloInstruction::CreateReshape(
-          ShapeUtil::MakeShape(lhs_shape.element_type(), lhs_reshape_dims,
-                               lhs_reshape_dynamic_dims),
-          transposed_lhs),
+      HloInstruction::CreateReshape(sh_lhs, transposed_lhs),
       &transposed_lhs->metadata());
 
   const auto& rhs_shape = original_dot->operand(1)->shape();
@@ -145,17 +164,29 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) {
   rhs_non_contracting_dims.reserve(num_rhs_non_contracting_dims);
   int64_t rhs_non_contracting_size = 1;
   bool rhs_non_contracting_dynamic = false;
+  int64_t rhs_non_contracting_multiplier_accu = 1;
+  DExpr rhs_non_contracting_expression = DExpr::Const(1);
   int64_t rhs_contracting_size = 1;
   bool rhs_contracting_dynamic = false;
+  int64_t rhs_contracting_multiplier_accu = 1;
+  DExpr rhs_contracting_expression = DExpr::Const(1);
+
+  bool rhs_contracting_is_static = true;
+  bool rhs_non_contracting_is_static = true;
+
   for (int64_t i = 0; i < rhs_rank; ++i) {
     if (absl::c_linear_search(original_dnums.rhs_contracting_dimensions(), i)) {
       rhs_contracting_size *= rhs_shape.dimensions(i);
       rhs_contracting_dynamic |= rhs_shape.is_dynamic_dimension(i);
+      rhs_contracting_expression =
+          rhs_contracting_expression * rhs_shape.expressions(i);
     } else if (!absl::c_linear_search(original_dnums.rhs_batch_dimensions(),
                                       i)) {
       rhs_non_contracting_dims.push_back(i);
       rhs_non_contracting_size *= rhs_shape.dimensions(i);
       rhs_non_contracting_dynamic |= rhs_shape.is_dynamic_dimension(i);
+      rhs_non_contracting_expression =
+          rhs_non_contracting_expression * rhs_shape.expressions(i);
     }
   }
 
@@ -184,27 +215,35 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) {
   rhs_reshape_dims.push_back(rhs_contracting_size);
   std::vector<bool> rhs_reshape_dynamic_dims = batch_dynamic_dims;
   rhs_reshape_dynamic_dims.push_back(rhs_contracting_dynamic);
+  std::vector<DExpr> rhs_reshape_expressions = batch_expressions;
+  rhs_reshape_expressions.push_back(rhs_contracting_expression);
   if (rhs_non_contracting_size > 1) {
     rhs_reshape_dims.push_back(rhs_non_contracting_size);
     rhs_reshape_dynamic_dims.push_back(rhs_non_contracting_dynamic);
+    rhs_reshape_expressions.push_back(rhs_non_contracting_expression);
   }
   // Reshape the contracting and non-contracting dimensions together.
+  auto sh_rhs = ShapeUtil::MakeShape(rhs_shape.element_type(), rhs_reshape_dims,
+                                     rhs_reshape_dynamic_dims,
+                                     rhs_reshape_expressions);
   HloInstruction* reshaped_rhs = computation->AddInstruction(
       HloInstruction::CreateReshape(
-          ShapeUtil::MakeShape(rhs_shape.element_type(), rhs_reshape_dims,
-                               rhs_reshape_dynamic_dims),
+          sh_rhs,
           transposed_rhs),
       &transposed_rhs->metadata());
 
   std::vector<int64_t> dot_dims = batch_dim_sizes;
   std::vector<bool> dot_dynamic_dims = batch_dynamic_dims;
+  std::vector<DExpr> dot_expressions = batch_expressions;
   if (lhs_non_contracting_size > 1) {
     dot_dims.push_back(lhs_non_contracting_size);
     dot_dynamic_dims.push_back(lhs_non_contracting_dynamic);
+    dot_expressions.push_back(lhs_non_contracting_expression);
   }
   if (rhs_non_contracting_size > 1) {
     dot_dims.push_back(rhs_non_contracting_size);
     dot_dynamic_dims.push_back(rhs_non_contracting_dynamic);
+    dot_expressions.push_back(rhs_non_contracting_expression);
   }
 
   DotDimensionNumbers dot_dnums;
@@ -251,12 +290,13 @@ absl::Status CanonicalizeDot(HloDotInstruction* original_dot) {
         HloInstruction::CreateReshape(result_shape, meta), &meta->metadata());
     sparse_meta.push_back(meta);
   }
+  auto sh_dot =
+      ShapeUtil::MakeShape(original_dot->shape().element_type(), dot_dims,
+                           dot_dynamic_dims, dot_expressions);
 
   HloInstruction* dot = computation->AddInstruction(HloInstruction::CreateDot(
-      ShapeUtil::MakeShape(original_dot->shape().element_type(), dot_dims,
-                           dot_dynamic_dims),
-      reshaped_lhs, reshaped_rhs, dot_dnums, original_dot->precision_config(),
-      sparsity, sparse_meta));
+      sh_dot, reshaped_lhs, reshaped_rhs, dot_dnums,
+      original_dot->precision_config(), sparsity, sparse_meta));
   original_dot->SetupDerivedInstruction(dot);
 
   std::unique_ptr<HloInstruction> replacement =
diff --git a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc
index 33752d60cae8ce..6ddf023ffd77a7 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc
@@ -158,8 +158,10 @@ void ApplyJacobiRotationOverRows(Eigh2x2 rotation, XlaOp& tl, XlaOp& tr,
   Shape shape = tl.builder()->GetShape(tl).value();
   std::vector<int64_t> broadcast_dims(shape.dimensions().size() - 1);
   absl::c_iota(broadcast_dims, 0);
-  auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims);
-  auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims);
+  auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims,
+                          shape.expressions());
+  auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims,
+                          shape.expressions());
 
   auto s_conj = MaybeConjugate(s, true);
   std::tie(tl, tr, bl, br) =
@@ -179,8 +181,10 @@ void ApplyJacobiRotationOverCols(Eigh2x2 rotation, XlaOp& tl, XlaOp& tr,
   std::vector<int64_t> broadcast_dims(shape.dimensions().size() - 1);
   absl::c_iota(broadcast_dims, 0);
   broadcast_dims.back() = shape.dimensions().size() - 1;
-  auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims);
-  auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims);
+  auto c = BroadcastInDim(rotation.c, shape.dimensions(), broadcast_dims,
+                          shape.expressions());
+  auto s = BroadcastInDim(rotation.s, shape.dimensions(), broadcast_dims,
+                          shape.expressions());
 
   auto s_conj = MaybeConjugate(s, true);
   std::tie(tl, tr, bl, br) =
@@ -365,11 +369,12 @@ absl::Status EighExpander::SortByEigenvalues(XlaOp& v, XlaOp& w) {
   TF_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(w));
   const int64_t num_dims = v_shape.dimensions().size();
   auto dimensions = v_shape.dimensions();
+  auto expressions = v_shape.expressions();
 
   std::vector<int64_t> broadcast_dims(num_dims - 1);
   std::iota(broadcast_dims.begin(), broadcast_dims.end(), 0);
   broadcast_dims[num_dims - 2] = num_dims - 1;
-  w = BroadcastInDim(w, dimensions, broadcast_dims);
+  w = BroadcastInDim(w, dimensions, broadcast_dims, expressions);
 
   XlaOp sort_result =
       Sort({w, v},
diff --git a/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc
index dcf329f7c6d8b9..17775e0f132dfe 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/qr_expander.cc
@@ -58,6 +58,15 @@ std::vector<int64_t> ConcatVectors(absl::Span<const int64_t> xs,
   return output;
 }
 
+std::vector<DExpr> ConcatEVectors(absl::Span<const DExpr> xs,
+                                  absl::Span<const DExpr> ys) {
+  std::vector<DExpr> output;
+  output.reserve(xs.size() + ys.size());
+  std::copy(xs.begin(), xs.end(), std::back_inserter(output));
+  std::copy(ys.begin(), ys.end(), std::back_inserter(output));
+  return output;
+}
+
 // Computes sqrt(x^2 + y^2 + ...), avoiding overflow/underflow.
 // e.g. for 3 arguments:
 // def norm(x, y, z):
@@ -220,11 +229,15 @@ absl::StatusOr<QrDecomposition> QrExpander::QrBlock(
 
   const int64_t m = ShapeUtil::GetDimension(a_shape, -2);
   const int64_t n = ShapeUtil::GetDimension(a_shape, -1);
+  const DExpr& m_exp = ShapeUtil::GetExpression(a_shape, -2);
+  const DExpr& n_exp = ShapeUtil::GetExpression(a_shape, -1);
 
   const int64_t num_batch_dims = num_dims - 2;
   std::vector<int64_t> batch_dims(num_batch_dims);
+  std::vector<DExpr> batch_exprs(num_batch_dims);
   for (int i = 0; i < num_batch_dims; ++i) {
     batch_dims[i] = ShapeUtil::GetDimension(a_shape, i);
+    batch_exprs[i] = ShapeUtil::GetExpression(a_shape, i);
   }
 
   std::vector<int64_t> batch_dim_indices(num_batch_dims);
@@ -248,9 +261,12 @@ absl::StatusOr<QrDecomposition> QrExpander::QrBlock(
         minor_dim + 1);
 
     std::vector<int64_t> shape = batch_dims;
+    std::vector<DExpr> exprs = batch_exprs;
     shape.push_back(1);
     shape.push_back(m);
-    auto v_broadcast = Reshape(v, shape);
+    exprs.push_back(DExpr::Const(1));
+    exprs.push_back(m_exp);
+    auto v_broadcast = Reshape(v, shape, exprs);
     // a[:, j+1:] -= np.conj(tau) * (v[:, np.newaxis] @
     //     (np.conj(v[np.newaxis, :]) @ a[:, j+1:]))
     // We use masking rather than a loop-variant shape to handle the j+1:
@@ -263,7 +279,8 @@ absl::StatusOr<QrDecomposition> QrExpander::QrBlock(
 
     // a[j, j] = beta
     // a[j+1:,j] = v[j+1:]
-    auto iota = Reshape(Iota(a.builder(), S32, m), {m, 1});
+    auto iota =
+        Reshape(Iota(a.builder(), S32, m), {m, 1}, {m_exp, DExpr::Const(1)});
     auto predecessor_mask = ConvertElementType(Lt(iota, j), type);
     auto mask = Broadcast(ConvertElementType(Eq(iota, j), type),
                           std::vector<int64_t>(batch_dims.size(), 1));
@@ -279,7 +296,8 @@ absl::StatusOr<QrDecomposition> QrExpander::QrBlock(
     std::vector<int64_t> dim_ids(num_dims);
     std::iota(dim_ids.begin(), dim_ids.end(), 0);
     new_x = BroadcastInDim(new_x, ConcatVectors(batch_dims, {m, n}),
-                           /*broadcast_dimensions=*/dim_ids);
+                           /*broadcast_dimensions=*/dim_ids,
+                           ConcatEVectors(batch_exprs, {m_exp, n_exp}));
     a = Select(Eq(iota_mn, j), new_x, a);
 
     // taus[j] = tau
diff --git a/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc b/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc
index 2fe502429287b4..ca7eebb9cfb7e8 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/reduce_decomposer.cc
@@ -47,6 +47,7 @@ class VariadicReductionLayoutEqualizer : public DfsHloRewriteVisitor {
       if (first_input_s.layout() != input_s.layout()) {
         Shape new_input_s = ShapeUtil::MakeShapeWithDenseLayout(
             input_s.element_type(), input_s.dimensions(),
+            input_s.expressions(),
             first_input_s.layout().minor_to_major());
         auto copy = MakeCopyHlo(input, new_input_s);
         changed = true;
diff --git a/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc
index d41aaf49be5d32..2e84d07f076078 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/rng_bit_generator_expander.cc
@@ -66,7 +66,7 @@ RngBitGeneratorExpander::GetGeneratorComputation(const Shape& data_shape,
 
   XlaBuilder builder("rng");
   XlaOp state_param = Parameter(&builder, 0, state_shape, "state");
-  XlaOp key_op = Reshape(Slice(state_param, {0}, {1}, {1}), {});
+  XlaOp key_op = Reshape(Slice(state_param, {0}, {1}, {1}), {}, {});
   RngOutput output;
   switch (algorithm) {
     case RandomAlgorithm::RNG_THREE_FRY:
@@ -83,8 +83,8 @@ RngBitGeneratorExpander::GetGeneratorComputation(const Shape& data_shape,
                            RandomAlgorithm_Name(algorithm));
   }
 
-  XlaOp final_state =
-      ConcatInDim(&builder, {Reshape(key_op, {1}), output.state}, 0);
+  XlaOp final_state = ConcatInDim(
+      &builder, {Reshape(key_op, {1}, {DExpr::Const(1)}), output.state}, 0);
   Tuple(&builder, {final_state, output.value});
   TF_ASSIGN_OR_RETURN(XlaComputation xla_computation, builder.Build());
   TF_ASSIGN_OR_RETURN(HloComputation * new_computation,
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
index b3f1e108ba3122..814a0d2388f457 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
@@ -944,6 +944,7 @@ cc_library(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/time",
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 9fcdf7ce1eb7b7..171c1bdca2e2bf 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -87,6 +87,22 @@ namespace m = match;
 
 using primitive_util::NativeTypeOf;
 
+static bool HasDynamicConstantFrontendAttributes(
+    const HloInstruction* instruction);
+
+static bool HasDynamicConstantContents(const HloInstruction* instruction) {
+  if (!instruction->has_contents()) {
+    return false;
+  }
+  for (const auto& content : instruction->contents()) {
+    if (content.node_type_case() != ExpressionProto::kConstantValue &&
+        content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET) {
+      return true;
+    }
+  }
+  return false;
+}
+
 // Unwraps broadcasts hunting for a constant.  If we find one, checks if the
 // constant contains only the given value.
 bool IsAll(const HloInstruction* op, int8_t value) {
@@ -94,6 +110,9 @@ bool IsAll(const HloInstruction* op, int8_t value) {
     case HloOpcode::kBroadcast:
       return IsAll(op->operand(0), value);
     case HloOpcode::kConstant:
+      if (HasDynamicConstantFrontendAttributes(op)) {
+        return false;
+      }
       return op->literal().IsAll(value);
     default:
       return false;
@@ -107,6 +126,9 @@ bool IsAllFloat(const HloInstruction* op, float value) {
     case HloOpcode::kBroadcast:
       return IsAllFloat(op->operand(0), value);
     case HloOpcode::kConstant:
+      if (HasDynamicConstantFrontendAttributes(op)) {
+        return false;
+      }
       return op->literal().IsAllFloat(value);
     default:
       return false;
@@ -119,6 +141,9 @@ bool IsAll(const HloInstruction* op, const Literal& scalar) {
     case HloOpcode::kBroadcast:
       return IsAll(op->operand(0), scalar);
     case HloOpcode::kConstant:
+      if (HasDynamicConstantFrontendAttributes(op)) {
+        return false;
+      }
       return op->literal().IsAll(scalar);
     default:
       return false;
@@ -171,9 +196,18 @@ bool IsPositive(const HloInstruction* hlo,
 
 static bool IsScalarConstant(const HloInstruction* hlo) {
   return hlo->opcode() == HloOpcode::kConstant &&
+         !HasDynamicConstantFrontendAttributes(hlo) &&
          ShapeUtil::IsEffectiveScalar(hlo->shape());
 }
 
+static bool HasDynamicConstantFrontendAttributes(
+    const HloInstruction* instruction) {
+  const auto& attrs = instruction->frontend_attributes().map();
+  return HasDynamicConstantContents(instruction) ||
+         attrs.contains("dynamic_constant_index") ||
+         attrs.contains("dynamic_constant_expr");
+}
+
 std::optional<double> GetConstantValue(const HloInstruction* inst) {
   if (!IsScalarConstant(inst)) {
     return std::nullopt;
@@ -701,8 +735,9 @@ absl::Status AlgebraicSimplifierVisitor::ScalarMultiplyReduction(
     HloInstruction* multiplier;
     // Pattern match a scalar multiply.
     if (Match(inst, m::MultiplyAnyOrder(
-                        m::Op(&operand),
-                        m::Broadcast(m::ConstantScalar(&multiplier))))) {
+                        m::Op(&operand), m::Broadcast(m::ConstantScalar(
+                                            &multiplier)))) &&
+        !HasDynamicConstantFrontendAttributes(multiplier)) {
       CHECK_LT(index, user->operand_count());
       CHECK_EQ(inst, user->operands()[index]);
 
@@ -732,8 +767,9 @@ absl::Status AlgebraicSimplifierVisitor::ScalarMultiplyReduction(
     HloInstruction* operand;
     HloInstruction* multiplier;
     if (Match(inst, m::MultiplyAnyOrder(
-                        m::Op(&operand),
-                        m::Broadcast(m::ConstantScalar(&multiplier))))) {
+                        m::Op(&operand), m::Broadcast(m::ConstantScalar(
+                                            &multiplier)))) &&
+        !HasDynamicConstantFrontendAttributes(multiplier)) {
       values.push_back(*GetConstantValue(multiplier));
 
       TF_RETURN_IF_ERROR(inst->ReplaceAllUsesWith(operand));
@@ -903,6 +939,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
       Match(add, m::Add(m::Add(m::NonConstant(&a),
                                m::Broadcast(m::ConstantScalar(&c1))),
                         m::Broadcast(m::ConstantScalar(&c2))))) {
+    if (HasDynamicConstantFrontendAttributes(c1) ||
+        HasDynamicConstantFrontendAttributes(c2)) {
+      return absl::OkStatus();
+    }
     TF_ASSIGN_OR_RETURN(auto* sum_of_constants,
                         MakeBinaryHlo(HloOpcode::kAdd, c1, c2));
     if (ShapeUtil::IsScalar(sum_of_constants->shape()) &&
@@ -921,6 +961,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
       Match(add, m::Add(m::Subtract(m::Broadcast(m::ConstantScalar(&c1)),
                                     m::NonConstant(&a)),
                         m::Broadcast(m::ConstantScalar(&c2))))) {
+    if (HasDynamicConstantFrontendAttributes(c1) ||
+        HasDynamicConstantFrontendAttributes(c2)) {
+      return absl::OkStatus();
+    }
     TF_ASSIGN_OR_RETURN(HloInstruction * sum_of_constants,
                         MakeBinaryHlo(HloOpcode::kAdd, c1, c2));
     if (ShapeUtil::IsScalar(sum_of_constants->shape()) &&
@@ -1201,13 +1245,15 @@ absl::StatusOr<bool> AlgebraicSimplifierVisitor::TrySimplifyTautologicalCompare(
         m::Shape().IsEffectiveScalar().WithElementType(PrimitiveType::S32);
     if (Match(cmp, m::Compare(m::Op(&lhs),
                               m::Constant(&rhs).WithShape(scalar_shape_matcher))
-                       .WithComparisonDirection(ComparisonDirection::kLt))) {
+                       .WithComparisonDirection(ComparisonDirection::kLt)) &&
+        !HasDynamicConstantFrontendAttributes(rhs)) {
       return {LessThanCompareInfo{lhs, *rhs->literal().GetFirstInteger()}};
     } else if (Match(
                    cmp,
                    m::Compare(m::Constant(&lhs).WithShape(scalar_shape_matcher),
                               m::Op(&rhs))
-                       .WithComparisonDirection(ComparisonDirection::kGt))) {
+                       .WithComparisonDirection(ComparisonDirection::kGt)) &&
+               !HasDynamicConstantFrontendAttributes(lhs)) {
       return {LessThanCompareInfo{rhs, *lhs->literal().GetFirstInteger()}};
     }
     return std::nullopt;
@@ -1232,6 +1278,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAllGather(
   if (all_gather->shape().IsArray() &&
       Match(all_gather->mutable_operand(0),
             m::Broadcast(m::ConstantScalar()))) {
+    if (HasDynamicConstantFrontendAttributes(
+            all_gather->mutable_operand(0)->operand(0))) {
+      return absl::OkStatus();
+    }
     return ReplaceWithNewInstruction(
         all_gather,
         all_gather->mutable_operand(0)->CloneWithNewShape(all_gather->shape()));
@@ -1244,6 +1294,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleAllToAll(
   if (all_to_all->shape().IsArray() &&
       Match(all_to_all->mutable_operand(0),
             m::Broadcast(m::ConstantScalar()))) {
+    if (HasDynamicConstantFrontendAttributes(
+            all_to_all->mutable_operand(0)->operand(0))) {
+      return absl::OkStatus();
+    }
     return ReplaceInstruction(all_to_all, all_to_all->mutable_operand(0));
   }
   return absl::OkStatus();
@@ -2185,6 +2239,10 @@ std::unique_ptr<HloInstruction> TryDivideToShift(
     return nullptr;
   }
 
+  if (HasDynamicConstantFrontendAttributes(c)) {
+    return nullptr;
+  }
+
   if (ShapeUtil::ElementIsSigned(divide->shape())) {
     int64_t b_value = static_cast<int64_t>(c->literal().GetFirstElement<T>());
     if (b_value > 0 && absl::has_single_bit(static_cast<uint64_t>(b_value))) {
@@ -2234,6 +2292,13 @@ absl::Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) {
   CHECK(Match(divide, m::Divide(m::Op(&a), m::Op(&b))));
   // A/1 => A
   VLOG(10) << "trying transform [A/1 => A]: " << divide->ToString();
+  if ((b->opcode() == HloOpcode::kConstant &&
+       HasDynamicConstantFrontendAttributes(b)) ||
+      (b->opcode() == HloOpcode::kBroadcast &&
+       b->operand(0)->opcode() == HloOpcode::kConstant &&
+       HasDynamicConstantFrontendAttributes(b->operand(0)))) {
+    return absl::OkStatus();
+  }
   if (IsAll(b, 1) && ReplaceInstructionIfCompatible(divide, a)) {
     return absl::OkStatus();
   }
@@ -2317,6 +2382,29 @@ absl::Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) {
     return absl::OkStatus();
   }
 
+  // A / broadcast(B) => A * broadcast(1 / B).
+  //
+  // This changes floating-point rounding and the handling of exceptional
+  // values, so backends must explicitly opt in when reciprocal substitution
+  // is permitted.
+  if (options_.enable_divide_by_broadcast_reciprocal() &&
+      ShapeUtil::ElementIsFloating(divide->shape()) &&
+      Match(divide,
+            m::Divide(
+                m::Op(&a),
+                m::Broadcast(m::Op(&b).WithShape(m::Shape().IsScalar()))))) {
+    TF_ASSIGN_OR_RETURN(
+        HloInstruction * reciprocal,
+        MakeBinaryHlo(HloOpcode::kDivide, MakeScalarLike(b, 1), b));
+    HloInstruction* reciprocal_broadcast =
+        divide->mutable_operand(1)->AddInstruction(
+            HloInstruction::CreateBroadcast(divide->shape(), reciprocal, {}));
+    return ReplaceWithNewInstruction(
+        divide, HloInstruction::CreateBinary(divide->shape(),
+                                             HloOpcode::kMultiply, a,
+                                             reciprocal_broadcast));
+  }
+
   // A / Const => A * (1 / Const)
   //
   // (Backends can do this transformation, but generally only if the constant is
@@ -4453,6 +4541,10 @@ absl::StatusOr<std::unique_ptr<HloInstruction>> MinMaxToClamp(
   CHECK(Match(clamp_upper_bound_bcast,
               m::Broadcast(m::ConstantEffectiveScalar(&clamp_upper_bound))))
       << clamp_upper_bound_bcast->ToString();
+  if (HasDynamicConstantFrontendAttributes(clamp_lower_bound) ||
+      HasDynamicConstantFrontendAttributes(clamp_upper_bound)) {
+    return nullptr;
+  }
 
   const Literal& lower_bound =
       Cast<HloConstantInstruction>(clamp_lower_bound)->literal();
@@ -4697,6 +4789,9 @@ absl::Status AlgebraicSimplifierVisitor::HandleClamp(HloInstruction* clamp) {
   if ((Match(to_clamp, m::PartitionId()) || Match(to_clamp, m::ReplicaId())) &&
       Match(clamp_lower_bound, m::ConstantScalar(0U)) &&
       Match(clamp_upper_bound, m::ConstantScalar())) {
+    if (HasDynamicConstantFrontendAttributes(clamp_upper_bound)) {
+      return absl::OkStatus();
+    }
     int64_t upper_bound = Cast<HloConstantInstruction>(clamp_upper_bound)
                               ->literal()
                               .GetFirstElement<uint32_t>();
@@ -4924,6 +5019,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleMultiply(
             m::Multiply(
                 m::Multiply(m::Op(&a), m::Broadcast(m::ConstantScalar(&c1))),
                 m::Broadcast(m::ConstantScalar(&c2))))) {
+    if (HasDynamicConstantFrontendAttributes(c1) ||
+        HasDynamicConstantFrontendAttributes(c2)) {
+      return absl::OkStatus();
+    }
     TF_ASSIGN_OR_RETURN(auto* product_of_constants,
                         MakeBinaryHlo(HloOpcode::kMultiply, c1, c2));
     if (ShapeUtil::IsScalar(product_of_constants->shape()) &&
@@ -5975,6 +6074,9 @@ std::unique_ptr<HloInstruction> TryRemainderToAnd(
       !Match(b, m::Broadcast(m::ConstantEffectiveScalar(&c)))) {
     return nullptr;
   }
+  if (HasDynamicConstantFrontendAttributes(c)) {
+    return nullptr;
+  }
 
   if (ShapeUtil::ElementIsSigned(remainder->shape())) {
     int64_t b_value = static_cast<int64_t>(c->literal().GetFirstElement<T>());
@@ -6060,7 +6162,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleRemainder(
   HloInstruction* divisor;
   if (Match(remainder,
             m::Remainder(m::Iota(&iota),
-                         m::Broadcast(m::ConstantEffectiveScalar(&divisor))))) {
+                         m::Broadcast(m::ConstantEffectiveScalar(&divisor)))) &&
+      !HasDynamicConstantFrontendAttributes(divisor)) {
     // The iota counts {0, ..., iota_upper_bound - 1}.  (Actually this is
     // conservative; the iota may overflow and count up to a smaller value than
     // this.  But that's OK for our purposes here.)
@@ -6088,7 +6191,7 @@ absl::Status AlgebraicSimplifierVisitor::HandleRemainder(
               m::AddAnyOrder(m::Iota(&iota),
                              m::Broadcast(m::ConstantEffectiveScalar(&addend))),
               m::Broadcast(&bcast, m::ConstantEffectiveScalar(&divisor)))) &&
-      addend == divisor) {
+      addend == divisor && !HasDynamicConstantFrontendAttributes(divisor)) {
     // The iota counts {0, ...iota_upper_bound - 1}, with the same caveat above
     // that iota_upper_bound is conservative, and the true upper bound may be
     // smaller.
@@ -7309,6 +7412,9 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicSlice(
     std::vector<int64_t> slice_strides(rank, 1);
 
     for (int64_t i = 0; i < rank; ++i) {
+      if (HasDynamicConstantFrontendAttributes(dynamic_slice->operand(i + 1))) {
+        return absl::OkStatus();
+      }
       std::optional<int64_t> offset =
           dynamic_slice->operand(i + 1)->literal().GetFirstInteger();
       if (!offset || *offset < 0) {
@@ -7502,6 +7608,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice(
           compatible = false;
           break;
         }
+        if (HasDynamicConstantFrontendAttributes(slice_dim_start)) {
+          compatible = false;
+          break;
+        }
         VLOG(2) << "slice: " << slice_dim_start->ToString();
         std::optional<int64_t> start =
             slice_dim_start->literal().GetFirstInteger();
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h
index 1a2cab861649e3..a0f44014e5d712 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h
@@ -314,6 +314,15 @@ class AlgebraicSimplifierOptions {
   }
   bool enable_fast_math() const { return enable_fast_math_; }
 
+  // If true, replace elementwise division by a broadcast scalar with
+  // multiplication by a broadcast scalar reciprocal.
+  void set_enable_divide_by_broadcast_reciprocal(bool enable) {
+    enable_divide_by_broadcast_reciprocal_ = enable;
+  }
+  bool enable_divide_by_broadcast_reciprocal() const {
+    return enable_divide_by_broadcast_reciprocal_;
+  }
+
   void set_enable_broadcast_degenerate_dimension(
       bool enable_broadcast_degenerate_dimension) {
     enable_broadcast_degenerate_dimension_ =
@@ -388,6 +397,7 @@ class AlgebraicSimplifierOptions {
   bool use_convert_constant_folding_{false};
   bool disable_dynamic_slice_to_slice_conversion_{false};
   bool enable_fast_math_{false};
+  bool enable_divide_by_broadcast_reciprocal_{false};
   bool enable_broadcast_degenerate_dimension_{true};
   bool enable_remove_no_op_reduce_precision_{false};
   bool enable_onednn_support_{
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index 7cff2113e4aceb..14991fb8dca9cf 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -2016,6 +2016,61 @@ TEST_F(AlgebraicSimplifierTest, DivideByBroadcastedConstant) {
                   m::Broadcast(m::Op().IsConstantScalar(1.0f / 256.0f)))));
 }
 
+TEST_F(AlgebraicSimplifierTest,
+       DivideByBroadcastedRuntimeScalarUsingReciprocal) {
+  const char* kModuleStr = R"(
+    HloModule m
+    test {
+      p0 = f32[16] parameter(0)
+      p1 = f32[] parameter(1)
+      b = f32[16] broadcast(p1), dimensions={}
+      ROOT d = f32[16] divide(p0, b)
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+
+  ASSERT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::Divide(m::Parameter(0),
+                                   m::Broadcast(m::Parameter(1)))));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_divide_by_broadcast_reciprocal(true);
+  ASSERT_TRUE(AlgebraicSimplifier(options).Run(m.get()).value());
+  EXPECT_THAT(
+      m->entry_computation()->root_instruction(),
+      GmockMatch(m::Multiply(
+          m::Parameter(0),
+          m::Broadcast(m::Divide(m::ConstantScalar(1), m::Parameter(1))))));
+}
+
+TEST_F(AlgebraicSimplifierTest,
+       DivideByBroadcastedRuntimeScalarPreservesResultLayout) {
+  const char* kModuleStr = R"(
+    HloModule m
+    test {
+      p0 = f32[2,3]{1,0} parameter(0)
+      p1 = f32[] parameter(1)
+      b = f32[2,3]{1,0} broadcast(p1), dimensions={}
+      ROOT d = f32[2,3]{0,1} divide(p0, b)
+    }
+  )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_divide_by_broadcast_reciprocal(true);
+  ASSERT_TRUE(AlgebraicSimplifier(options).Run(m.get()).value());
+
+  HloInstruction* root = m->entry_computation()->root_instruction();
+  EXPECT_THAT(
+      root,
+      GmockMatch(m::Multiply(
+          m::Parameter(0),
+          m::Broadcast(m::Divide(m::ConstantScalar(1), m::Parameter(1))))));
+  EXPECT_TRUE(LayoutUtil::Equal(root->shape().layout(),
+                                LayoutUtil::MakeLayout({0, 1})));
+}
+
 // pow(pow(A, X), Y) => pow(A, X*Y)
 TEST_F(AlgebraicSimplifierTest, PowerOfPower) {
   auto m = CreateNewVerifiedModule();
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc
index 5a9ef91900d850..7ed0be00ab1d44 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/time.h"
 #include "xla/hlo/evaluator/hlo_evaluator.h"
@@ -225,6 +226,56 @@ absl::StatusOr<bool> HloConstantFolding::Run(
         continue;
       }
 
+      bool source_has_dynamic_constant_marker = false;
+      std::vector<std::string> marked_constant_operands;
+      if (instruction->has_contents()) {
+        bool has_dynamic_content = false;
+        for (const auto& content : instruction->contents()) {
+          has_dynamic_content =
+              content.node_type_case() != ExpressionProto::kConstantValue &&
+              content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET;
+          if (has_dynamic_content) {
+            break;
+          }
+        }
+        if (has_dynamic_content) {
+          source_has_dynamic_constant_marker = true;
+          marked_constant_operands.push_back(absl::StrFormat(
+              "%s:contents=%d", instruction->name(),
+              instruction->contents().size()));
+        }
+      }
+      for (const HloInstruction* operand : instruction->operands()) {
+        if (operand->opcode() != HloOpcode::kConstant ||
+            !operand->has_contents()) {
+          continue;
+        }
+        bool has_dynamic_content = false;
+        for (const auto& content : operand->contents()) {
+          has_dynamic_content =
+              content.node_type_case() != ExpressionProto::kConstantValue &&
+              content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET;
+          if (has_dynamic_content) {
+            break;
+          }
+        }
+        if (!has_dynamic_content) {
+          continue;
+        }
+        source_has_dynamic_constant_marker = true;
+        marked_constant_operands.push_back(absl::StrFormat(
+            "%s:contents=%d literal=%s", operand->name(),
+            operand->contents().size(), operand->literal().ToString()));
+      }
+      if (source_has_dynamic_constant_marker) {
+        VLOG(1) << "Skipping HloConstantFolding for " << instruction->name()
+                << " (" << HloOpcodeString(instruction->opcode())
+                << ") because it or its source constant operands carry dynamic contents";
+        VLOG(1) << "Marked constant operands: "
+                << absl::StrJoin(marked_constant_operands, ", ");
+        continue;
+      }
+
       // Check for instructions that we can't fold even if they appear inside of
       // a subcomputation (e.g. a kCall).
       if (IsOrContainsIllegalInstr(instruction)) {
@@ -323,6 +374,17 @@ absl::StatusOr<bool> HloConstantFolding::Run(
       changed = true;
       HloInstruction* new_constant = instruction->AddInstruction(
           HloInstruction::CreateConstant(std::move(result)));
+      VLOG(1) << "HloConstantFolding created constant from "
+              << instruction->name() << " ("
+              << HloOpcodeString(instruction->opcode())
+              << "), source_has_dynamic_constant_marker="
+              << source_has_dynamic_constant_marker;
+      if (!marked_constant_operands.empty()) {
+        VLOG(1) << "Marked constant operands: "
+                << absl::StrJoin(marked_constant_operands, ", ");
+      }
+      VLOG(1) << "Folded constant literal -> "
+              << new_constant->literal().ToString();
       if (new_constant->shape().has_layout()) {
         // Update element_size_in_bits on the new instruction's layout. Literals
         // always have element_size_in_bits set to 0, and CreateConstant copies
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc
index ea348d312eaaa6..d8e03d1f98de09 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/optimize_input_output_buffer_alias_test.cc
@@ -40,9 +40,11 @@ class OptimizeInputOutputBufferAliasTest
     r2f32_ = ShapeUtil::MakeShape(F32, {4, 5});
     r3f32_ = ShapeUtil::MakeShape(F32, {4, 5, 6});
     r4f32_ = ShapeUtil::MakeShape(F32, {4, 5, 6, 7});
-    d1f32_ = ShapeUtil::MakeShape(F32, {256}, /*dynamic_dimensions=*/{true});
+    d1f32_ = ShapeUtil::MakeShape(F32, {256}, /*dynamic_dimensions=*/{true},
+                                  /*expressions=*/{});
     d2f32_ = ShapeUtil::MakeShape(F32, {128, 128},
-                                  /*dynamic_dimensions=*/{false, true});
+                                  /*dynamic_dimensions=*/{false, true},
+                                  /*expressions=*/{});
     // Static shape with same size as dynamic shape `d1f32_`.
     d3f32_ = ShapeUtil::MakeShape(F32, {512});
   }
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
index 8391636e00daac..6e50c950e1d7e4 100644
--- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
+++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
@@ -2673,6 +2673,8 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) {
   SmallVector<xla::XlaOp> dimSizes;
   SmallVector<int64_t> newSizeBounds;
   std::vector<bool> dimsAreDynamic;
+  std::vector<xla::DExpr> dimExpressions;
+
   for (auto i = 0; i < resultType.getRank(); ++i) {
     auto runtimeSizeX1 = xla::Slice(outputShape, {i}, {i + 1}, {1});
     dimSizes.push_back(xla::Reshape(runtimeSizeX1, {}));
@@ -2683,9 +2685,10 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) {
       return op->emitOpError() << "unbounded dynamism is not supported";
     newSizeBounds.push_back(hlo::isStaticDimSize(dimSize) ? dimSize : dimBound);
     dimsAreDynamic.push_back(!hlo::isStaticDimSize(dimSize));
+    dimExpressions.push_back(xla::DExpr::Unknown(40));
   }
-  value_map[op] =
-      xla::DynamicReshape(operand, dimSizes, newSizeBounds, dimsAreDynamic);
+  value_map[op] = xla::DynamicReshape(operand, dimSizes, newSizeBounds,
+                                      dimsAreDynamic, dimExpressions);
   return success();
 }
 
@@ -2696,7 +2699,8 @@ LogicalResult ExportXlaOp(ReshapeOp op, OpLoweringContext ctx) {
     return failure();
 
   value_map[op] =
-      xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions());
+      xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions(),
+                   xla::TypeToShape(op.getType()).expressions());
   return success();
 }
 
@@ -3004,6 +3008,7 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) {
   SmallVector<xla::XlaOp> dimSizes;
   SmallVector<int64_t> newSizeBounds;
   std::vector<bool> dimsAreDynamic;
+  std::vector<xla::DExpr> dimExpressions;
   for (auto i = 0; i < resultType.getRank(); ++i) {
     auto runtimeSizeX1 = xla::Slice(outputShape, {i}, {i + 1}, {1});
     dimSizes.push_back(xla::Reshape(runtimeSizeX1, {}));
@@ -3014,9 +3019,11 @@ LogicalResult ExportXlaOp(DynamicReshapeOp op, OpLoweringContext ctx) {
       return op->emitOpError() << "unbounded dynamism is not supported";
     newSizeBounds.push_back(hlo::isStaticDimSize(dimSize) ? dimSize : dimBound);
     dimsAreDynamic.push_back(!hlo::isStaticDimSize(dimSize));
+    dimExpressions.push_back(xla::DExpr::Unknown(50));
   }
   value_map[op] =
-      xla::DynamicReshape(operand, dimSizes, newSizeBounds, dimsAreDynamic);
+      xla::DynamicReshape(operand, dimSizes, newSizeBounds, dimsAreDynamic,
+                          dimExpressions);
   return success();
 }
 
@@ -4558,7 +4565,8 @@ LogicalResult ExportXlaOp(ReshapeOp op, OpLoweringContext ctx) {
     return failure();
 
   value_map[op] =
-      xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions());
+      xla::Reshape(operand, xla::TypeToShape(op.getType()).dimensions(),
+                   xla::TypeToShape(op.getType()).expressions());
   return success();
 }
 
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc
index 7bfbbff39956ce..a06d517914bdf6 100644
--- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc
+++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/type_to_shape.cc
@@ -150,6 +150,7 @@ Shape TypeToShape(mlir::Type type) {
 
     llvm::SmallVector<int64_t, 4> shape(rank, mlir::ShapedType::kDynamic);
     std::vector<bool> is_dynamic(rank, false);
+    std::vector<DExpr> expressions(rank, DExpr::Unknown(60));
     for (int64_t dim = 0; dim < rank; ++dim) {
       int64_t size = t.getDimSize(dim);
       if (size == ShapedType::kDynamic) {
@@ -191,7 +192,8 @@ Shape TypeToShape(mlir::Type type) {
       return sparse_shape;
     }
 
-    return ShapeUtil::MakeShape(primitive_type, shape, is_dynamic);
+    return ShapeUtil::MakeShape(primitive_type, shape, is_dynamic,
+                                expressions);
   } else if (auto tuple_type = mlir::dyn_cast<mlir::TupleType>(type)) {
     llvm::SmallVector<Shape, 4> shapes;
     shapes.reserve(tuple_type.size());
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 7cd03e529ae86a..bc3afa6f73ba2a 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -121,6 +121,26 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "dynamic_constant_rewriter",
+    srcs = ["dynamic_constant_rewriter.cc"],
+    hdrs = ["dynamic_constant_rewriter.h"],
+    deps = [
+        "//xla:literal_util",
+        "//xla:shape_util",
+        "//xla:status_macros",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@local_tsl//tsl/platform:errors",
+        "@local_tsl//tsl/platform:protobuf",
+    ],
+)
+
 cc_library(
     name = "all_reduce_promotion",
     srcs = ["all_reduce_promotion.cc"],
diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 24247e60993937..8ea07dccc5a25e 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -295,6 +295,7 @@ cc_library(
         "//xla/mlir_hlo:all_passes",
         "//xla/mlir_hlo:transforms_passes",
         "//xla/service:all_reduce_promotion",
+        "//xla/service:dynamic_constant_rewriter",
         "//xla/service:all_to_all_decomposer",
         "//xla/service:batched_gather_scatter_normalizer",
         "//xla/service:batchnorm_expander",
@@ -597,6 +598,7 @@ cc_library(
         ":onednn_layer_norm",
         ":onednn_matmul",
         ":onednn_softmax",
+        ":runtime_batch_size_api",
         ":runtime_conv2d",
         ":runtime_conv2d_acl",
         ":runtime_conv2d_mkl",
@@ -689,6 +691,7 @@ cc_library(
         "//xla/service:custom_call_status_internal",
         "//xla/service:executable",
         "//xla/service:hlo_execution_profile",
+        "//xla/service:hlo_profile_printer",
         "//xla/service:hlo_profile_printer_data_cc",
         "//xla/service:hlo_value",
         "//xla/service:maybe_owning_device_memory",
@@ -1094,6 +1097,7 @@ cc_library(
     name = "cpu_runtime",
     srcs = [
         "cpu_runtime.cc",
+        "runtime_batch_size.cc",
         "xfeed_manager.cc",
     ],
     hdrs = [
@@ -1103,6 +1107,7 @@ cc_library(
     copts = runtime_copts(),
     deps = [
         ":cpu_executable_run_options",
+        ":runtime_batch_size_api",
         "//xla:executable_run_options",
         "//xla:shape_util",
         "//xla:types",
@@ -1122,6 +1127,7 @@ cc_library(
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/concurrency:async_value",
         "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:env_time",
         "//xla/tsl/platform:logging",
         "//xla/tsl/platform:status",
         "@com_google_absl//absl/algorithm:container",
@@ -1156,6 +1162,11 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "runtime_batch_size_api",
+    hdrs = ["runtime_batch_size.h"],
+)
+
 cc_library(
     name = "runtime_conv3d",
     srcs = ["runtime_conv3d.cc"],
@@ -1461,6 +1472,7 @@ xla_cc_test(
     tags = ["optonly"],
     deps = [
         ":cpu_runtime",
+        ":runtime_batch_size_api",
         ":runtime_custom_call_status",
         ":runtime_matmul",
         ":runtime_matmul_acl",
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc
index fcffdafcf5e5ee..5d04d299cecfc2 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc
@@ -40,6 +40,7 @@ limitations under the License.
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
@@ -152,6 +153,7 @@ limitations under the License.
 #include "xla/map_util.h"
 #include "xla/mlir_hlo/transforms/passes.h"
 #include "xla/service/all_reduce_promotion.h"
+#include "xla/service/dynamic_constant_rewriter.h"
 #include "xla/service/all_to_all_decomposer.h"
 #include "xla/service/batched_gather_scatter_normalizer.h"
 #include "xla/service/batchnorm_expander.h"
@@ -450,6 +452,23 @@ void AddHloVerifier(HloPassPipeline* pipeline, HloVerifierOpts&& opts = {},
   }
 }
 
+bool DynamicConstantRewriterEnabled(const HloModule* module) {
+  const auto& extra_options =
+      module->config().debug_options().xla_backend_extra_options();
+  auto it = extra_options.find("xla_cpu_disable_dynamic_constant_rewriter");
+  if (it == extra_options.end()) {
+    return true;
+  }
+  bool disabled = false;
+  if (!absl::SimpleAtob(it->second, &disabled)) {
+    LOG(WARNING)
+        << "Ignoring invalid xla_cpu_disable_dynamic_constant_rewriter value: "
+        << it->second;
+    return true;
+  }
+  return !disabled;
+}
+
 std::unique_ptr<HloPassFix<HloPassPipeline>> CreateSimplificationPipeline(
     absl::string_view name, HloModule* module, bool is_fusion_emitters) {
   // Run the following passes to a fixed point.
@@ -465,6 +484,11 @@ std::unique_ptr<HloPassFix<HloPassPipeline>> CreateSimplificationPipeline(
       !module->config().debug_options().xla_cpu_enable_fast_min_max());
   options.set_supports_non_canonical_dots(false);
   options.set_executing_on_cpu(true);
+  options.set_enable_divide_by_broadcast_reciprocal(
+      module->config().debug_options().xla_cpu_enable_fast_math() &&
+      !module->config()
+           .debug_options()
+           .xla_cpu_fast_math_honor_division());
   pipeline->AddPass<AlgebraicSimplifier>(options);
   pipeline->AddPass<SortSimplifier>();
   pipeline->AddPass<HloDCE>();
@@ -475,7 +499,7 @@ std::unique_ptr<HloPassFix<HloPassPipeline>> CreateSimplificationPipeline(
   }
 
   // Needs to happen after algebraic simplifier.
-  pipeline->AddPass<TreeReductionRewriter>();
+  // pipeline->AddPass<TreeReductionRewriter>();
 
   // BatchNormExpander can create zero-sized ops, so zero-sized HLO
   // elimination has to come after that pass.
@@ -495,6 +519,9 @@ std::unique_ptr<HloPassFix<HloPassPipeline>> CreateSimplificationPipeline(
       options::FoldAllConstants(module->config())
           ? HloConstantFolding::Level::kAggressive
           : HloConstantFolding::Level::kDefault);
+  if (DynamicConstantRewriterEnabled(module)) {
+    pipeline->AddPass<DynamicConstantRewriter>();
+  }
   pipeline->AddPass<ConditionalSimplifier>();
 
   return pipeline;
@@ -885,6 +912,11 @@ absl::Status CpuCompiler::RunHloPassesAfterLayoutAssn(
     options.set_minmax_propagate_nan(
         !module->config().debug_options().xla_cpu_enable_fast_min_max());
     options.set_executing_on_cpu(true);
+    options.set_enable_divide_by_broadcast_reciprocal(
+        module->config().debug_options().xla_cpu_enable_fast_math() &&
+        !module->config()
+             .debug_options()
+             .xla_cpu_fast_math_honor_division());
     pipeline.AddPass<AlgebraicSimplifier>(options);
     pipeline.AddPass<HloDCE>();
     pipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/true);
@@ -1041,6 +1073,25 @@ absl::Status CreateHloProfilingArtifacts(
   return absl::OkStatus();
 }
 
+bool ShouldCreateHloProfilingArtifacts(const HloModule& module) {
+  if (!module.config().hlo_profiling_enabled()) {
+    return false;
+  }
+
+  // The thunk runtime launches host kernels through XLA_CPU_KernelCallFrame,
+  // which currently has no profile-counters field. Emitting profiling loads and
+  // stores in this mode can make generated code dereference a null counter
+  // pointer at run time.
+  if (module.config().debug_options().xla_cpu_use_thunk_runtime()) {
+    LOG(WARNING) << "--xla_hlo_profile is not supported by XLA:CPU thunk "
+                    "runtime; disabling HLO profiling for module "
+                 << module.name();
+    return false;
+  }
+
+  return true;
+}
+
 }  // namespace
 
 absl::StatusOr<std::unique_ptr<HloModule>> CpuCompiler::RunHloPasses(
@@ -1429,7 +1480,7 @@ CpuCompiler::CompileCpuExecutable(std::unique_ptr<HloModule> module) {
       computation_to_profile_idx;
   std::unique_ptr<HloProfileIndexMap> hlo_profile_index_map;
   std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data;
-  if (module->config().hlo_profiling_enabled()) {
+  if (ShouldCreateHloProfilingArtifacts(*module)) {
     TF_RETURN_IF_ERROR(CreateHloProfilingArtifacts(
         *module, &instruction_to_profile_idx, &computation_to_profile_idx,
         &hlo_profile_index_map, &hlo_profile_printer_data));
@@ -1981,7 +2032,7 @@ CpuCompiler::CompileAheadOfTimeLegacy(
   std::unique_ptr<HloProfileIndexMap> hlo_profile_index_map;
   std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data;
 
-  if (module->config().hlo_profiling_enabled()) {
+  if (ShouldCreateHloProfilingArtifacts(*module)) {
     TF_RETURN_IF_ERROR(CreateHloProfilingArtifacts(
         *module, &instruction_to_profile_idx, &computation_to_profile_idx,
         &hlo_profile_index_map, &hlo_profile_printer_data));
@@ -2139,7 +2190,7 @@ CpuCompiler::CompileAheadOfTimeThunks(
       computation_to_profile_idx;
   std::unique_ptr<HloProfileIndexMap> hlo_profile_index_map;
   std::unique_ptr<HloProfilePrinterData> hlo_profile_printer_data;
-  if (module->config().hlo_profiling_enabled()) {
+  if (ShouldCreateHloProfilingArtifacts(*module)) {
     TF_RETURN_IF_ERROR(CreateHloProfilingArtifacts(
         *module, &instruction_to_profile_idx, &computation_to_profile_idx,
         &hlo_profile_index_map, &hlo_profile_printer_data));
@@ -2388,7 +2439,7 @@ CpuCompiler::CompileAheadOfTimeThunks(
       cpu_executable->thunks().thunk_sequence();
 
   std::unique_ptr<HloProfilePrinterData> executable_hlo_profile_printer_data =
-      cpu_executable->module().config().hlo_profiling_enabled()
+      cpu_executable->hlo_profiling_enabled()
           ? std::make_unique<HloProfilePrinterData>(
                 cpu_executable->hlo_profile_printer_data())
           : nullptr;
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc
index cefc0274106501..11e9e35a511cd7 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc
@@ -27,7 +27,10 @@ limitations under the License.
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Casting.h"
 #include "xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_config.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/service/llvm_compiler.h"
 #include "xla/tests/hlo_test_base.h"
@@ -164,6 +167,45 @@ TEST_F(CpuCompilerInternalsTest, JustOneDylibWithThunks) {
                          << "one dylib is allowed.";
 }
 
+TEST_F(CpuCompilerInternalsTest, SharesBroadcastScalarReciprocal) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
+                          ParseAndReturnVerifiedModule(R"(
+    HloModule test
+    ENTRY main {
+      p0 = f32[16] parameter(0)
+      p1 = f32[16] parameter(1)
+      scalar = f32[] parameter(2)
+      broadcast = f32[16] broadcast(scalar), dimensions={}
+      divide0 = f32[16] divide(p0, broadcast)
+      divide1 = f32[16] divide(p1, broadcast)
+      ROOT result = f32[16] add(divide0, divide1)
+    }
+  )"));
+  DebugOptions& debug_options =
+      hlo_module->mutable_config().mutable_debug_options();
+  debug_options.set_xla_cpu_enable_fast_math(true);
+  debug_options.set_xla_cpu_fast_math_honor_division(false);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> optimized_module,
+                          GetOptimizedModule(std::move(hlo_module)));
+
+  int64_t divide_count = 0;
+  int64_t scalar_divide_count = 0;
+  for (const HloComputation* computation : optimized_module->computations()) {
+    for (const HloInstruction* instruction : computation->instructions()) {
+      if (instruction->opcode() != HloOpcode::kDivide) {
+        continue;
+      }
+      ++divide_count;
+      if (instruction->shape().dimensions().empty()) {
+        ++scalar_divide_count;
+      }
+    }
+  }
+  EXPECT_EQ(divide_count, 1);
+  EXPECT_EQ(scalar_divide_count, 1);
+}
+
 }  // namespace
 }  // namespace cpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/cpu/cpu_executable.cc b/third_party/xla/xla/service/cpu/cpu_executable.cc
index 158a4000d27722..0307561329ee27 100644
--- a/third_party/xla/xla/service/cpu/cpu_executable.cc
+++ b/third_party/xla/xla/service/cpu/cpu_executable.cc
@@ -54,6 +54,7 @@ limitations under the License.
 #include "xla/service/custom_call_status_internal.h"
 #include "xla/service/executable.h"
 #include "xla/service/hlo_execution_profile.h"
+#include "xla/service/hlo_profile_printer.h"
 #include "xla/service/hlo_profile_printer_data.pb.h"
 #include "xla/service/hlo_value.h"
 #include "xla/service/maybe_owning_device_memory.h"
@@ -76,6 +77,20 @@ limitations under the License.
 namespace xla {
 namespace cpu {
 
+namespace {
+
+float GetClockRateGhz(const ExecutableRunOptions* run_options) {
+  if (run_options->stream() == nullptr ||
+      run_options->stream()->parent() == nullptr) {
+    return 1.0f;
+  }
+  float clock_rate_ghz =
+      run_options->stream()->parent()->GetDeviceDescription().clock_rate_ghz();
+  return clock_rate_ghz > 0 ? clock_rate_ghz : 1.0f;
+}
+
+}  // namespace
+
 absl::StatusOr<std::unique_ptr<CpuExecutable>> CpuExecutable::Create(
     std::unique_ptr<FunctionLibrary> function_library,
     std::unique_ptr<const BufferAssignment> assignment,
@@ -163,14 +178,34 @@ static absl::StatusOr<MaybeOwningDeviceMemory> MemoryForAllocation(
     se::DeviceMemoryAllocator* memory_allocator, int device_ordinal) {
   VLOG(3) << allocation.ToString();
   if (allocation.is_entry_computation_parameter()) {
-    se::DeviceMemoryBase out = arguments[allocation.parameter_number()]
+    se::DeviceMemoryBase param_mem = arguments[allocation.parameter_number()]
                                    .Buffer(allocation.param_shape_index())
                                    .AsDeviceMemoryBase();
-    CHECK_LE(allocation.size(), out.size())
-        << "Size mismatch on param " << allocation.parameter_number()
-        << " at shape index " << allocation.param_shape_index().ToString();
-    VLOG(3) << "allocation is a parameter";
-    return MaybeOwningDeviceMemory{out};
+
+    const int64_t compiled_bytes = allocation.size();
+    const int64_t runtime_bytes = param_mem.size();
+
+    if(runtime_bytes == 0){
+      return MaybeOwningDeviceMemory{param_mem};
+    }
+
+    if (compiled_bytes <= runtime_bytes) {
+      return MaybeOwningDeviceMemory{param_mem};
+    }
+
+    //padded owns that device memory
+    TF_ASSIGN_OR_RETURN(
+        se::OwningDeviceMemory padded,
+        memory_allocator->Allocate(device_ordinal, compiled_bytes));
+
+    void* dst = padded->opaque();
+    void* src = param_mem.opaque();
+
+    std::memcpy(dst, src, runtime_bytes);
+    //Fill rest of them with zeros.
+    std::memset(static_cast<char*>(dst) + runtime_bytes, 0, compiled_bytes - runtime_bytes);
+    return MaybeOwningDeviceMemory{std::move(padded)};
+
   } else if (allocation.is_constant()) {
     VLOG(3) << "allocation is a constant";
     if (allocation.index() < constants.size()) {
@@ -226,8 +261,16 @@ absl::Status CpuExecutable::ExecuteComputeFunction(
     absl::Span<MaybeOwningDeviceMemory const> buffers) {
   uint64_t start_micros = tsl::Env::Default()->NowMicros();
 
-  size_t profile_counters_size = 0;
+  std::vector<int64_t> profile_counter_storage;
+  if (hlo_profiling_enabled()) {
+    profile_counter_storage.assign(
+        hlo_profile_printer_data().profile_counters_size(), 0);
+  }
+  size_t profile_counters_size = profile_counter_storage.size();
   int64_t* profile_counters = nullptr;
+  if (!profile_counter_storage.empty()) {
+    profile_counters = profile_counter_storage.data();
+  }
 
   // Call the computation function following the calling convention. See the
   // definition of 'ComputeFunctionType' for the details of the calling
@@ -266,6 +309,15 @@ absl::Status CpuExecutable::ExecuteComputeFunction(
   compute_function_(nullptr, run_options, nullptr, buffer_pointers.data(),
                     &status, profile_counters);
   record_profile();
+  if (profile_counters != nullptr) {
+    std::string hlo_profile =
+        PrintHloProfile(hlo_profile_printer_data(), profile_counters,
+                        GetClockRateGhz(run_options));
+    if (!hlo_profile.empty()) {
+      LOG(INFO) << "XLA:CPU HLO profile for " << module().name() << "\n"
+                << hlo_profile;
+    }
+  }
   std::optional<absl::string_view> error_message =
       CustomCallStatusGetMessage(&status);
   if (error_message) {
@@ -275,11 +327,28 @@ absl::Status CpuExecutable::ExecuteComputeFunction(
   return absl::OkStatus();
 }
 
+void PrintScalars(absl::Span<MaybeOwningDeviceMemory const> buffers) {
+  for (int i = 0; i < buffers.size(); ++i) {
+    const se::DeviceMemoryBase& dmem = buffers[i].AsDeviceMemoryBase();
+    if (dmem.opaque() && dmem.size() >= sizeof(int64_t)) {
+      int64_t val = 0;
+      std::memcpy(&val, dmem.opaque(), sizeof(val));
+      std::cerr << "Buffer " << i << " scalar: " << val << std::endl;
+    } else {
+      std::cerr << "Buffer " << i << " empty or too small" << std::endl;
+    }
+  }
+}
+
 absl::Status CpuExecutable::ExecuteThunks(
     const ExecutableRunOptions* run_options,
     absl::Span<MaybeOwningDeviceMemory const> buffers) {
   uint64_t start_ns = tsl::Env::Default()->NowNanos();
 
+  #if defined(PRINT_BATCHSIZE)
+  PrintScalars(buffers);
+  #endif
+
   size_t profile_counters_size = 0;
   int64_t* profile_counters = nullptr;
 
@@ -318,7 +387,8 @@ absl::Status CpuExecutable::ExecuteThunks(
       intra_op_thread_pool,
       &task_runner,
       &collective_execute_params,
-      &custom_call_execute_params};
+      &custom_call_execute_params,
+      run_options->batch_size()};
 
   auto executed_event = thunks_->Execute(execute_params);
   tsl::BlockUntilReady(executed_event);
@@ -339,13 +409,15 @@ absl::StatusOr<ExecutionOutput> CpuExecutable::CreateResultShapedBuffer(
     absl::Span<MaybeOwningDeviceMemory> buffers,
     absl::Span<ExecutionInput> arguments) {
   se::Stream* stream = run_options->stream();
-  ExecutionOutput result(/*on_device_shape=*/result_shape(),
-                         run_options->allocator(),
-                         stream->parent()->device_ordinal());
   const HloInputOutputAliasConfig& input_output_alias =
       module().input_output_alias_config();
   HloInstruction* root = hlo_module_->entry_computation()->root_instruction();
   const Shape& root_shape = root->shape();
+  // Use root_shape to initialize ExecutionOuput as the batch multiplier info
+  // is only attached the ROOT
+  ExecutionOutput result(/*on_device_shape=*/root_shape, //result_shape(),
+                         run_options->allocator(),
+                         stream->parent()->device_ordinal());
 
   // Move se::OwningDeviceMemory values which contain the array(s) of the result
   // into the respective location in ScopedShapedBuffer which is returned to the
diff --git a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc
index bd26c1c87621cd..54c34104f4b7b0 100644
--- a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc
+++ b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc
@@ -59,6 +59,23 @@ bool IsNonComplexNonBatchedMatrixVectorDot(const HloInstruction* hlo) {
          hlo->dot_dimension_numbers().lhs_batch_dimensions_size() == 0;
 }
 
+bool HasDynamicDimensions(const Shape& shape) {
+  for (int64_t i = 0; i < shape.dimensions().size(); ++i) {
+    if (shape.is_dynamic_dimension(i) ||
+        shape.expressions(i)->is_dynamic()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool IsDynamicDot(const HloInstruction* hlo) {
+  return hlo->opcode() == HloOpcode::kDot &&
+         (HasDynamicDimensions(hlo->shape()) ||
+          HasDynamicDimensions(hlo->operand(0)->shape()) ||
+          HasDynamicDimensions(hlo->operand(1)->shape()));
+}
+
 bool HasExactlyOneUse(const HloInstruction& hlo_instr) {
   return hlo_instr.user_count() == 1 &&
          absl::c_count(hlo_instr.users().front()->operands(), &hlo_instr) == 1;
@@ -142,6 +159,10 @@ FusionDecision CpuInstructionFusion::ShouldFuse(HloInstruction* consumer,
     return FusionDecision::Forbid("Don't fuse large constants.");
   }
 
+  if (IsDynamicDot(producer) || IsDynamicDot(consumer)) {
+    return FusionDecision::Forbid("Do not fuse dynamic dots on CPU.");
+  }
+
   if (CanBeOutputFused(producer, consumer)) {
     VLOG(2) << "Fusion OK: Can create output fusion.";
     return FusionDecision::Allow();
diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.cc b/third_party/xla/xla/service/cpu/cpu_runtime.cc
index 7caf9c43b1119b..20c482f663b6c6 100644
--- a/third_party/xla/xla/service/cpu/cpu_runtime.cc
+++ b/third_party/xla/xla/service/cpu/cpu_runtime.cc
@@ -60,6 +60,7 @@ limitations under the License.
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/env_time.h"
 #include "xla/tsl/platform/logging.h"
 #include "xla/tsl/platform/status.h"
 #include "xla/util.h"
@@ -166,6 +167,8 @@ extern const char* const kParallelForkJoinSymbolName =
     "__xla_cpu_runtime_ParallelForkJoin";
 extern const char* const kPrintfToStderrSymbolName =
     "__xla_cpu_runtime_PrintfToStderr";
+extern const char* const kReadCycleCounterSymbolName =
+    "__xla_cpu_runtime_ReadCycleCounter";
 extern const char* const kStatusIsSuccessSymbolName =
     "__xla_cpu_runtime_StatusIsSuccess";
 extern const char* const kKeyValueSortSymbolName =
@@ -620,6 +623,11 @@ ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY int __xla_cpu_runtime_PrintfToStderr(
   return result;
 }
 
+ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY uint64_t
+__xla_cpu_runtime_ReadCycleCounter() {
+  return tsl::EnvTime::NowNanos();
+}
+
 ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY int64_t __xla_cpu_runtime_TracingStart(
     const void* /* ExecutableRunOptions*  run_options_ptr*/, const char* name,
     const char* hlo_module, int64_t program_id) {
diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.h b/third_party/xla/xla/service/cpu/cpu_runtime.h
index 71e27ea600ee28..de1ea93a855a35 100644
--- a/third_party/xla/xla/service/cpu/cpu_runtime.h
+++ b/third_party/xla/xla/service/cpu/cpu_runtime.h
@@ -79,6 +79,7 @@ extern const char* const kAcquireOutfeedBufferForPopulationSymbolName;
 extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName;
 extern const char* const kParallelForkJoinSymbolName;
 extern const char* const kPrintfToStderrSymbolName;
+extern const char* const kReadCycleCounterSymbolName;
 extern const char* const kStatusIsSuccessSymbolName;
 extern const char* const kKeyValueSortSymbolName;
 extern const char* const kTopKF32SymbolName;
@@ -115,6 +116,7 @@ int GetDeviceOrdinal(const xla::ExecutableRunOptions* run_options);
 extern "C" {
 
 extern int __xla_cpu_runtime_PrintfToStderr(const char* format, ...);
+extern uint64_t __xla_cpu_runtime_ReadCycleCounter();
 
 extern int64_t __xla_cpu_runtime_TracingStart(
     const void* /* xla::ExecutableRunOptions* */ run_options_ptr,
diff --git a/third_party/xla/xla/service/cpu/cpu_runtime_test.cc b/third_party/xla/xla/service/cpu/cpu_runtime_test.cc
index 79df0e29e8ea91..e14835602224a4 100644
--- a/third_party/xla/xla/service/cpu/cpu_runtime_test.cc
+++ b/third_party/xla/xla/service/cpu/cpu_runtime_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/array2d.h"
 #include "xla/client/local_client.h"
 #include "xla/executable_run_options.h"
+#include "xla/service/cpu/runtime_batch_size.h"
 #include "xla/service/cpu/runtime_custom_call_status.h"
 #include "xla/service/cpu/runtime_matmul.h"
 #include "xla/service/cpu/runtime_matmul_acl.h"
@@ -41,6 +42,14 @@ namespace {
 
 class CpuRuntimeTest : public ::testing::Test {};
 
+TEST_F(CpuRuntimeTest, GetBatchSize) {
+  ExecutableRunOptions run_options;
+  run_options.set_batch_size(42);
+
+  EXPECT_EQ(__xla_cpu_runtime_GetBatchSize(&run_options), 42);
+  EXPECT_EQ(__xla_cpu_runtime_GetBatchSize(nullptr), 0);
+}
+
 template <typename T>
 std::unique_ptr<Array2D<float>> MaybeTransposeArray2D(const Array2D<T>& array,
                                                       bool transpose) {
diff --git a/third_party/xla/xla/service/cpu/dot_op_emitter.cc b/third_party/xla/xla/service/cpu/dot_op_emitter.cc
index e2e4f914579dbb..8b877429123cbf 100644
--- a/third_party/xla/xla/service/cpu/dot_op_emitter.cc
+++ b/third_party/xla/xla/service/cpu/dot_op_emitter.cc
@@ -160,6 +160,24 @@ bool CanEmitTiledLlvmIrGemm(
   return true;
 }
 
+bool HasDynamicMatmulDims(const DotInfo& dot_info) {
+  const Shape& lhs_shape = dot_info.lhs_shape;
+  const Shape& rhs_shape = dot_info.rhs_shape;
+  const DotDimensionNumbers& dim_nums = dot_info.dim_nums;
+
+  DExpr m_expr = lhs_shape.dimensions().size() <= 1
+                     ? DExpr::Const(1)
+                     : lhs_shape.expressions(
+                           1LL - dim_nums.lhs_contracting_dimensions(0));
+  DExpr k_expr = lhs_shape.expressions(dim_nums.lhs_contracting_dimensions(0));
+  DExpr n_expr = rhs_shape.dimensions().size() <= 1
+                     ? DExpr::Const(1)
+                     : rhs_shape.expressions(
+                           1LL - dim_nums.rhs_contracting_dimensions(0));
+
+  return m_expr->is_dynamic() || k_expr->is_dynamic() || n_expr->is_dynamic();
+}
+
 // Returns dot implementation strategy for non-batch dot operations.
 DotImplementationStrategy GetNonBatchDotImplementationStrategy(
     const HloModuleConfig& config, const DotInfo& dot_info,
@@ -173,6 +191,10 @@ DotImplementationStrategy GetNonBatchDotImplementationStrategy(
          dot_info.dim_nums.rhs_batch_dimensions_size() == 0)
       << "Dot operations must be non-batch";
 
+  if (HasDynamicMatmulDims(dot_info)) {
+    return DotImplementationStrategy::kEigen;
+  }
+
   // Any Matrix-Vector product of floating point or integral type, or
   // a transpose-dot fusion of the same can be lowered to a tiled LLVM
   // IR implementation.
@@ -253,6 +275,12 @@ class DotOpEmitter {
     // The number of columns on the RHS.
     int64_t n;
 
+    DExpr m_expr;
+
+    DExpr k_expr;
+
+    DExpr n_expr;
+
     // True if the LHS matrix is column major.
     bool lhs_column_major;
 
@@ -858,16 +886,23 @@ absl::Status DotOpEmitter::EmitCallToRuntime() {
 
   if (!mat_mult_dims.lhs_column_major) {
     std::swap(mat_mult_dims.m, mat_mult_dims.n);
+    std::swap(mat_mult_dims.m_expr, mat_mult_dims.n_expr);
     std::swap(lhs, rhs);
     std::swap(transpose_lhs, transpose_rhs);
   }
 
-  b_->CreateCall(matmul_func,
-                 {executable_run_options_value_, target_array_.GetBasePointer(),
-                  lhs->GetBasePointer(), rhs->GetBasePointer(),
-                  b_->getInt64(mat_mult_dims.m), b_->getInt64(mat_mult_dims.n),
-                  b_->getInt64(mat_mult_dims.k), b_->getInt32(transpose_lhs),
-                  b_->getInt32(transpose_rhs)});
+  llvm::Value* m_val =
+      xla::llvm_ir::EmitExpression(b_, mat_mult_dims.m_expr);
+  llvm::Value* n_val =
+      xla::llvm_ir::EmitExpression(b_, mat_mult_dims.n_expr);
+  llvm::Value* k_val =
+      xla::llvm_ir::EmitExpression(b_, mat_mult_dims.k_expr);
+
+  b_->CreateCall(
+      matmul_func,
+      {executable_run_options_value_, target_array_.GetBasePointer(),
+       lhs->GetBasePointer(), rhs->GetBasePointer(), m_val, n_val, k_val,
+       b_->getInt32(transpose_lhs), b_->getInt32(transpose_rhs)});
   return absl::OkStatus();
 }
 
@@ -942,18 +977,30 @@ absl::Status DotOpEmitter::EmitCallToBatchRuntime() {
 
   if (!mat_mult_dims.lhs_column_major) {
     std::swap(mat_mult_dims.m, mat_mult_dims.n);
+    std::swap(mat_mult_dims.m_expr, mat_mult_dims.n_expr);
     std::swap(lhs, rhs);
     std::swap(transpose_lhs, transpose_rhs);
   }
 
+  llvm::Value* m_val =
+      xla::llvm_ir::EmitExpression(b_, mat_mult_dims.m_expr);
+  llvm::Value* n_val =
+      xla::llvm_ir::EmitExpression(b_, mat_mult_dims.n_expr);
+  llvm::Value* k_val =
+      xla::llvm_ir::EmitExpression(b_, mat_mult_dims.k_expr);
+  DExpr batch_size_expr =
+      lhs_shape.expressions(0) ? lhs_shape.expressions(0)
+                               : DExpr::Const(lhs_shape.dimensions(0));
+  llvm::Value* batch_size_val =
+      xla::llvm_ir::EmitExpression(b_, batch_size_expr);
+
   VLOG(1) << "Batch dot emitted with runtime:" << fn_name;
 
   b_->CreateCall(
       matmul_func,
       {executable_run_options_value_, target_array_.GetBasePointer(),
-       lhs->GetBasePointer(), rhs->GetBasePointer(),
-       b_->getInt64(mat_mult_dims.m), b_->getInt64(mat_mult_dims.n),
-       b_->getInt64(mat_mult_dims.k), b_->getInt64(lhs_shape.dimensions(0)),
+       lhs->GetBasePointer(), rhs->GetBasePointer(), m_val, n_val, k_val,
+       batch_size_val,
        b_->getInt32(static_cast<uint32_t>(transpose_lhs)),
        b_->getInt32(static_cast<uint32_t>(transpose_rhs))});
   return absl::OkStatus();
@@ -983,6 +1030,14 @@ DotOpEmitter::MatMultDims DotOpEmitter::GetMatMultDims() const {
       /*n=*/rhs_shape.dimensions().size() <= 1
           ? 1LL
           : rhs_shape.dimensions(1LL - dim_nums.rhs_contracting_dimensions(0)),
+      /*m_expr=*/lhs_shape.dimensions().size() <= 1
+          ? DExpr::Const(1)
+          : lhs_shape.expressions(1LL - dim_nums.lhs_contracting_dimensions(0)),
+      /*k_expr=*/
+      lhs_shape.expressions(dim_nums.lhs_contracting_dimensions(0)),
+      /*n_expr=*/rhs_shape.dimensions().size() <= 1
+          ? DExpr::Const(1)
+          : rhs_shape.expressions(1LL - dim_nums.rhs_contracting_dimensions(0)),
       /*lhs_column_major=*/is_column_major(lhs_shape),
       /*lhs_canonical=*/lhs_shape.dimensions().size() <= 1 ||
           dim_nums.lhs_contracting_dimensions(0) == 1,
@@ -1014,6 +1069,14 @@ DotOpEmitter::MatMultDims DotOpEmitter::GetBatchMatMultDims() const {
       /*n=*/rhs_shape.dimensions().size() <= 1
           ? 1LL
           : rhs_shape.dimensions(2LL - dim_nums.rhs_contracting_dimensions(0)),
+      /*m_expr=*/lhs_shape.dimensions().size() <= 1
+          ? DExpr::Const(1)
+          : lhs_shape.expressions(2LL - dim_nums.lhs_contracting_dimensions(0)),
+      /*k_expr=*/
+      lhs_shape.expressions(1LL + dim_nums.lhs_contracting_dimensions(0)),
+      /*n_expr=*/rhs_shape.dimensions().size() <= 1
+          ? DExpr::Const(1)
+          : rhs_shape.expressions(2LL - dim_nums.rhs_contracting_dimensions(0)),
       /*lhs_column_major=*/is_column_major(lhs_shape),
       /*lhs_canonical=*/lhs_shape.dimensions().size() <= 1 ||
           dim_nums.lhs_contracting_dimensions(0) == 1,
@@ -1093,22 +1156,34 @@ absl::Status EmitNonBatchDotOperation(
 
 Shape DropFirstDim(const Shape& shape) {
   absl::Span<int64_t const> array_shape_dims(shape.dimensions());
+  absl::Span<const DExpr> array_shape_exprs(shape.expressions());
   array_shape_dims.remove_prefix(1);
-  return ShapeUtil::MakeShapeWithDescendingLayout(shape.element_type(),
-                                                  array_shape_dims);
+  array_shape_exprs.remove_prefix(1);
+  return ShapeUtil::MakeShapeWithDescendingLayout(
+      shape.element_type(), array_shape_dims, array_shape_exprs);
 }
 
 Shape CollapseFirstNDims(const Shape& shape, int64_t n) {
   absl::Span<int64_t const> input_shape_dims(shape.dimensions());
+  absl::Span<const DExpr> input_expressions(shape.expressions());
   int64_t prefix_dim =
       std::accumulate(input_shape_dims.begin(), input_shape_dims.begin() + n,
                       1ll, std::multiplies<int64_t>());
+
+  DExpr prefix_expression = std::accumulate(
+      input_expressions.begin(), input_expressions.begin() + n, DExpr::Const(1),
+      [](DExpr acc, const DExpr& v) { return acc * v; });
+
   DimensionVector result_dims;
+  std::vector<DExpr> result_expressions;
   result_dims.push_back(prefix_dim);
+  result_expressions.push_back(prefix_expression);
   std::copy(input_shape_dims.begin() + n, input_shape_dims.end(),
             std::back_inserter(result_dims));
-  return ShapeUtil::MakeShapeWithDescendingLayout(shape.element_type(),
-                                                  result_dims);
+  std::copy(input_expressions.begin() + n, input_expressions.end(),
+            std::back_inserter(result_expressions));
+  return ShapeUtil::MakeShapeWithDescendingLayout(
+      shape.element_type(), result_dims, result_expressions);
 }
 
 llvm_ir::IrArray CollapseFirstNDims(llvm::IRBuilderBase* b,
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc
index feca6552d243f8..bdf1bc2e60f4df 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/ir_emitter.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "xla/shape_expr.h"
 #include "absl/meta/type_traits.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -76,6 +77,7 @@ limitations under the License.
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/map_util.h"
+#include "xla/printer.h"
 #include "xla/primitive_util.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/collective_ops_utils.h"
@@ -2055,9 +2057,11 @@ absl::Status IrEmitter::HandleSlice(HloInstruction* slice) {
   }
 
   llvm_ir::IrArray source_array = GetIrArrayFor(operand);
+  const auto* slice_hlo = Cast<HloSliceInstruction>(slice);
   const llvm_ir::IrArray::Index source_index = target_index.SourceIndexOfSlice(
-      /*operand_shape=*/operand->shape(), /*starts=*/slice->slice_starts(),
-      /*strides=*/slice->slice_strides(), /*builder=*/b());
+      /*operand_shape=*/operand->shape(), /*starts=*/slice_hlo->slice_starts(),
+      /*start_exprs=*/slice_hlo->slice_start_exprs(),
+      /*strides=*/slice_hlo->slice_strides(), /*builder=*/b());
 
   llvm::Value* memcpy_dest =
       target_array.EmitArrayElementAddress(target_index, b(), "slice.dest");
@@ -2067,7 +2071,7 @@ absl::Status IrEmitter::HandleSlice(HloInstruction* slice) {
   const int64_t memcpy_elements =
       primitive_elements_per_logical_element * memcpy_logical_elements;
 
-  EmitTransferElements(memcpy_dest, memcpy_source, memcpy_elements,
+  EmitTransferElements(memcpy_dest, memcpy_source, DExpr::Const(memcpy_elements),
                        slice->shape().element_type(), target_array,
                        source_array);
 
@@ -2358,6 +2362,26 @@ absl::Status IrEmitter::HandleSliceToDynamic(HloInstruction* hlo) {
   return EmitSliceToDynamic(hlo, source_arrays, target_array);
 }
 
+absl::Status IrEmitter::HandleShapeExprValue(HloInstruction* hlo) {
+  TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
+
+  llvm_ir::IrArray out_array = GetIrArrayFor(hlo);
+  TF_RET_CHECK(hlo->has_contents());
+  TF_RET_CHECK(!hlo->contents().empty());
+  xla::DExpr expr = xla::DExprFromProto(hlo->contents()[0]);
+  TF_RET_CHECK(expr);
+  llvm::Value* expr_value = llvm_ir::EmitExpression(b(), expr);
+
+  auto it = emitted_value_.find(hlo);
+  if (it == emitted_value_.end()) {
+    LOG(ERROR) << "No buffer assigned for instruction " << hlo->name();
+  }
+  llvm::Value* dest_ptr = it->second;
+  b()->CreateStore(expr_value, dest_ptr);
+
+  return absl::OkStatus();
+}
+
 absl::Status IrEmitter::HandlePadToStatic(HloInstruction* hlo) {
   TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
 
@@ -2806,6 +2830,10 @@ absl::Status IrEmitter::HandleOneDnnSoftmax(HloInstruction* custom_call) {
 #endif  // INTEL_MKL
 
 absl::Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) {
+  if (custom_call->custom_call_target() == "GetExpressionValue") {
+    return HandleShapeExprValue(custom_call);
+  }
+
   if (custom_call->custom_call_target() == "PadToStatic") {
     return HandlePadToStatic(custom_call);
   }
@@ -3125,11 +3153,11 @@ absl::Status EmitFastConcatenate(
   // contiguous subregion in the target buffer starting at target_region_begin.
   llvm::Value* target_region_begin =
       target_array.EmitArrayElementAddress(target_index, &b, "target_region");
-  int64_t byte_offset_into_target_region = 0;
+  llvm::Value* byte_offset_into_target_region = b.getInt64(0);
 
-  int64_t inner_dims_product = absl::c_accumulate(
-      inner_dims, int64_t{1}, [&](int64_t product, int64_t inner_dim) {
-        return product * output_shape.dimensions(inner_dim);
+  DExpr inner_exprs_product = absl::c_accumulate(
+      inner_dims, DExpr::Const(1), [&](DExpr product, int64_t inner_dim) {
+        return product * output_shape.expressions(inner_dim);
       });
 
   // For each operand, emit a memcpy from the operand to the target of size
@@ -3142,18 +3170,24 @@ absl::Status EmitFastConcatenate(
     llvm::Value* copy_source_address =
         source_array.EmitArrayElementAddress(source_index, &b, "src_addr");
 
-    llvm::Value* copy_target_address =
-        b.CreateGEP(b.getInt8Ty(), target_region_begin,
-                    b.getInt64(byte_offset_into_target_region));
+    llvm::Value* copy_target_address = b.CreateGEP(
+        b.getInt8Ty(), target_region_begin, byte_offset_into_target_region);
+
+    auto cexpr = input_shape.expressions(concat_dim);
 
     ::xla::cpu::EmitTransferElements(
         copy_target_address, copy_source_address,
-        inner_dims_product * input_shape.dimensions(concat_dim), primitive_type,
-        target_array, source_array, module, b);
+        inner_exprs_product * cexpr, primitive_type, target_array,
+        source_array, module, b);
+
+    llvm::Value* concat_dim_count = xla::llvm_ir::EmitExpression(
+        &b, inner_exprs_product * input_shape.expressions(concat_dim));
 
-    byte_offset_into_target_region += inner_dims_product *
-                                      input_shape.dimensions(concat_dim) *
-                                      primitive_type_size;
+    llvm::Value* concat_dim_size =
+        b.CreateMul(concat_dim_count, b.getInt64(primitive_type_size));
+    byte_offset_into_target_region =
+        b.CreateAdd(byte_offset_into_target_region, concat_dim_size,
+                    "byte_offset_into_target_region");
   }
 
   if (!outer_dims.empty()) {
@@ -3364,7 +3398,7 @@ llvm::Value* IrEmitter::EmitCallToFfi(HloCustomCallInstruction* custom_call,
 }
 
 void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source,
-                                     int64_t element_count,
+                                     const xla::DExpr& element_count,
                                      PrimitiveType primitive_type,
                                      const llvm_ir::IrArray& target_array,
                                      const llvm_ir::IrArray& source_array) {
@@ -3374,7 +3408,8 @@ void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source,
 }
 
 void EmitTransferElements(llvm::Value* target, llvm::Value* source,
-                          int64_t element_count, PrimitiveType primitive_type,
+                          const xla::DExpr& element_count,
+                          PrimitiveType primitive_type,
                           const llvm_ir::IrArray& target_array,
                           const llvm_ir::IrArray& source_array,
                           llvm::Module* module, llvm::IRBuilderBase& b) {
@@ -3386,7 +3421,7 @@ void EmitTransferElements(llvm::Value* target, llvm::Value* source,
   llvm::Type* primitive_llvm_type =
       llvm_ir::PrimitiveTypeToIrType(primitive_type, module->getContext());
 
-  if (element_count == 1) {
+  if (element_count->is_constant() && element_count->get_val() == 1) {
     auto* load_instruction =
         b.CreateAlignedLoad(primitive_llvm_type, source, element_alignment);
     source_array.AnnotateLoadStoreInstructionWithMetadata(load_instruction);
@@ -3394,11 +3429,12 @@ void EmitTransferElements(llvm::Value* target, llvm::Value* source,
         b.CreateAlignedStore(load_instruction, target, element_alignment);
     target_array.AnnotateLoadStoreInstructionWithMetadata(store_instruction);
   } else {
+    auto element_count_value = xla::llvm_ir::EmitExpression(&b, element_count);
+    llvm::Value* elements_size =
+        b.CreateMul(element_count_value, b.getInt64(primitive_type_size));
     auto* memcpy_instruction = b.CreateMemCpy(
         target, /*DstAlign=*/llvm::Align(element_alignment), source,
-        /*SrcAlign=*/llvm::Align(element_alignment),
-        element_count * primitive_type_size);
-
+        /*SrcAlign=*/llvm::Align(element_alignment), elements_size);
     // The memcpy does the load and the store internally.  The aliasing related
     // metadata has to reflect that.
     std::map<int, llvm::MDNode*> merged_metadata =
@@ -3679,10 +3715,26 @@ void IrEmitter::ProfilingState::UpdateProfileCounter(llvm::IRBuilderBase* b,
                                                      llvm::Value* prof_counter,
                                                      llvm::Value* cycle_end,
                                                      llvm::Value* cycle_start) {
+  llvm::Value* profile_counters = prof_counter;
+  if (auto* gep = llvm::dyn_cast<llvm::GetElementPtrInst>(prof_counter)) {
+    profile_counters = gep->getPointerOperand();
+  }
+
+  auto* profile_counters_type =
+      llvm::cast<llvm::PointerType>(profile_counters->getType());
+  llvm::Value* has_profile_counters = b->CreateICmpNE(
+      profile_counters, llvm::ConstantPointerNull::get(profile_counters_type),
+      "has_profile_counters");
+
+  llvm::AllocaInst* dummy_counter = llvm_ir::EmitAllocaAtFunctionEntry(
+      b->getInt64Ty(), "dummy_profile_counter", b);
+  b->CreateStore(b->getInt64(0), dummy_counter);
+  prof_counter =
+      b->CreateSelect(has_profile_counters, prof_counter, dummy_counter);
+
   auto* cycle_diff = b->CreateSub(cycle_end, cycle_start);
   llvm::LoadInst* old_cycle_count = b->CreateLoad(
-      llvm::cast<llvm::GetElementPtrInst>(prof_counter)->getSourceElementType(),
-      prof_counter, "old_cycle_count");
+      b->getInt64Ty(), prof_counter, "old_cycle_count");
   auto* new_cycle_count =
       b->CreateAdd(cycle_diff, old_cycle_count, "new_cycle_count");
   b->CreateStore(new_cycle_count, prof_counter);
@@ -3692,10 +3744,17 @@ llvm::Value* IrEmitter::ProfilingState::ReadCycleCounter(
     llvm::IRBuilderBase* b) {
   llvm::Module* module = b->GetInsertBlock()->getModule();
   if (!use_rdtscp_) {
-    llvm::Function* func_llvm_readcyclecounter =
-        llvm::Intrinsic::getOrInsertDeclaration(
-            module, llvm::Intrinsic::readcyclecounter);
-    return b->CreateCall(func_llvm_readcyclecounter);
+    llvm::FunctionType* fn_type =
+        llvm::FunctionType::get(b->getInt64Ty(), /*isVarArg=*/false);
+    llvm::FunctionCallee read_cycle_counter_func =
+        module->getOrInsertFunction(runtime::kReadCycleCounterSymbolName,
+                                    fn_type);
+    if (auto* fn =
+            llvm::dyn_cast<llvm::Function>(read_cycle_counter_func.getCallee())) {
+      fn->setCallingConv(llvm::CallingConv::C);
+      fn->setDoesNotThrow();
+    }
+    return b->CreateCall(read_cycle_counter_func);
   }
   llvm::Function* func_llvm_x86_rdtscp =
       llvm::Intrinsic::getOrInsertDeclaration(module,
@@ -3911,8 +3970,10 @@ llvm::Value* IrEmitter::EmitThreadLocalBufferPointer(
 
       if (!target_shape.IsOpaque()) {
         AttachAlignmentMetadataForLoad(param_address_untyped, target_shape);
-        AttachDereferenceableMetadataForLoad(param_address_untyped,
-                                             target_shape);
+        if (!target_shape.has_dynamic_expr()) {
+          AttachDereferenceableMetadataForLoad(param_address_untyped,
+                                               target_shape);
+        }
       }
       return param_address_untyped;
     }
@@ -3958,7 +4019,10 @@ llvm::Value* IrEmitter::EmitGlobalBufferPointer(
 
   AttachInvariantLoadMetadataForLoad(tempbuf_address_base);
   AttachAlignmentMetadataForLoad(tempbuf_address_base, allocation.size());
-  AttachDereferenceableMetadataForLoad(tempbuf_address_base, allocation.size());
+
+  if (!target_shape.has_dynamic_expr())
+    AttachDereferenceableMetadataForLoad(tempbuf_address_base,
+                                         allocation.size());
 
   llvm::Value* tempbuf_address_untyped = tempbuf_address_base;
   // Any explicit buffer pointer should point to the start of the slice.
@@ -4059,9 +4123,38 @@ absl::Status IrEmitter::EmitMemcpy(const HloInstruction& source,
   llvm::Value* source_value = GetEmittedValueFor(&source);
   llvm::Value* destination_value = GetEmittedValueFor(&destination);
   int64_t source_size = ByteSizeOf(source.shape());
-  // TODO(b/63762267): Be more aggressive about specifying alignment.
-  MemCpy(destination_value, /*DstAlign=*/llvm::Align(1), source_value,
-         /*SrcAlign=*/llvm::Align(1), source_size);
+  auto shape = source.shape();
+  auto expressions = shape.expressions();
+  bool is_dynamic =
+      std::any_of(expressions.begin(), expressions.end(),
+                  [](const DExpr& e) { return e && e->is_dynamic(); });
+  if (is_dynamic) {
+    llvm::LLVMContext& ctx = b()->getContext();
+    llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx);
+    int64_t dimensions_accu = 1;
+    DExpr expression_accu = DExpr::Const(1);
+    for (int i = 0; i < shape.dimensions_size(); i++) {
+      const auto& expression = shape.expressions(i);
+      if (expression && expression->is_dynamic()) {
+        dimensions_accu *= shape.dimensions(i);
+        expression_accu = expression_accu * expression;
+      }
+    }
+    llvm::Value* expr_value =
+        xla::llvm_ir::EmitExpression(b(), expression_accu);
+    // Divide the size in bytes by the size of the dynamic dimension(s).
+    // TODO: make that less hacky
+    llvm::ConstantInt* size =
+        llvm::ConstantInt::get(i64Type, source_size / dimensions_accu, true);
+    llvm::Value* memcopy_size =
+        b()->CreateMul(expr_value, size, "memcopy_size");
+    MemCpy(destination_value, /*DstAlign=*/llvm::Align(1), source_value,
+           /*SrcAlign=*/llvm::Align(1), memcopy_size);
+  } else {
+    // TODO(b/63762267): Be more aggressive about specifying alignment.
+    MemCpy(destination_value, /*DstAlign=*/llvm::Align(1), source_value,
+           /*SrcAlign=*/llvm::Align(1), source_size);
+  }
   return absl::OkStatus();
 }
 
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.h b/third_party/xla/xla/service/cpu/ir_emitter.h
index 40f54d2f4bff97..53583810acd5cb 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.h
+++ b/third_party/xla/xla/service/cpu/ir_emitter.h
@@ -332,6 +332,7 @@ class IrEmitter : public DfsHloVisitorWithDefault,
 
  private:
   absl::Status HandleSliceToDynamic(HloInstruction* hlo);
+  absl::Status HandleShapeExprValue(HloInstruction* hlo);
   absl::Status HandlePadToStatic(HloInstruction* hlo);
   absl::Status HandleTopK(HloInstruction* hlo) override;
   absl::Status HandleAllReduceSingleReplica(HloInstruction* crs);
@@ -569,7 +570,8 @@ class IrEmitter : public DfsHloVisitorWithDefault,
   // Emits LLVM IR to transfer "element_count" elements of type "primitive_type"
   // from the address "source" to the address "target".
   void EmitTransferElements(llvm::Value* target, llvm::Value* source,
-                            int64_t element_count, PrimitiveType primitive_type,
+                            const xla::DExpr& element_count,
+                            PrimitiveType primitive_type,
                             const llvm_ir::IrArray& target_array,
                             const llvm_ir::IrArray& source_array);
 
@@ -859,7 +861,8 @@ class IrEmitter : public DfsHloVisitorWithDefault,
 
 // Decoupled implementation of IrEmitter::EmitTransferElements.
 void EmitTransferElements(llvm::Value* target, llvm::Value* source,
-                          int64_t element_count, PrimitiveType primitive_type,
+                          const xla::DExpr& element_count,
+                          PrimitiveType primitive_type,
                           const llvm_ir::IrArray& target_array,
                           const llvm_ir::IrArray& source_array,
                           llvm::Module* module, llvm::IRBuilderBase& b);
diff --git a/third_party/xla/xla/service/cpu/ir_emitter2.cc b/third_party/xla/xla/service/cpu/ir_emitter2.cc
index bce2108bb87572..e39fba7a0b9098 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter2.cc
+++ b/third_party/xla/xla/service/cpu/ir_emitter2.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
+#include "xla/shape_expr.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
@@ -74,6 +75,7 @@ limitations under the License.
 #include "xla/service/llvm_ir/ir_array.h"
 #include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/service/llvm_ir/loop_emitter.h"
+#include "xla/printer.h"
 #include "xla/shape.h"
 #include "xla/shape_partition.h"
 #include "xla/stream_executor/launch_dim.h"
@@ -195,6 +197,30 @@ absl::StatusOr<IrEmitter2::KernelInfo> IrEmitter2::EmitPadHostKernel(
       KernelInfo(std::move(kernel_prototype), se::BlockDim(), se::ThreadDim()));
 }
 
+absl::StatusOr<IrEmitter2::KernelInfo>
+IrEmitter2::EmitGetExpressionValueHostKernel(const HloInstruction* getBatch) {
+  VLOG(2) << "Emit GetExpressionValue host kernel: " << getBatch->name();
+
+  TF_ASSIGN_OR_RETURN(KernelPrototype kernel_prototype,
+                      EmitKernelPrototype(getBatch));
+  llvm_ir::IrArray operand_array = kernel_prototype.arguments[0];
+  llvm_ir::IrArray output_array = kernel_prototype.results[0];
+  TF_RET_CHECK(getBatch->has_contents());
+  TF_RET_CHECK(!getBatch->contents().empty());
+  xla::DExpr expr = xla::DExprFromProto(getBatch->contents()[0]);
+  TF_RET_CHECK(expr);
+  llvm::IRBuilder<> b(module_->getContext());
+  b.SetInsertPoint(kernel_prototype.function->getEntryBlock().getTerminator());
+  llvm::Value* bdim_value = llvm_ir::EmitExpression(&b, expr);
+  llvm_ir::IrArray::Index output_index(/*multidimensional_index=*/{},
+                                       getBatch->shape(), b.getInt32Ty());
+  llvm::Value* output_ptr =
+      output_array.EmitArrayElementAddress(output_index, &b);
+  b.CreateStore(bdim_value, output_ptr);
+  return kernels_.emplace_back(
+      KernelInfo(std::move(kernel_prototype), se::BlockDim(), se::ThreadDim()));
+}
+
 absl::StatusOr<IrEmitter2::KernelInfo> IrEmitter2::EmitFusionHostKernel(
     const HloFusionInstruction* fusion) {
   VLOG(2) << "Emit fusion host kernel: " << fusion->name();
diff --git a/third_party/xla/xla/service/cpu/ir_emitter2.h b/third_party/xla/xla/service/cpu/ir_emitter2.h
index e720e06f37642c..8c493a164626cf 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter2.h
+++ b/third_party/xla/xla/service/cpu/ir_emitter2.h
@@ -113,6 +113,9 @@ class IrEmitter2 {
   // Emits a host kernel for the pad instruction.
   absl::StatusOr<KernelInfo> EmitPadHostKernel(const HloInstruction* pad);
 
+  absl::StatusOr<KernelInfo> EmitGetExpressionValueHostKernel(
+      const HloInstruction* getBatch);
+
   // Emits a host kernel for the given fusion instruction.
   absl::StatusOr<KernelInfo> EmitFusionHostKernel(
       const HloFusionInstruction* fusion);
diff --git a/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc b/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc
index 2bfffd88df937e..e3af0399381b79 100644
--- a/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc
+++ b/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc
@@ -61,6 +61,7 @@ ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock(absl::string_view loop_name,
     // performance with a large improvement in compile time.
     auto unroll_mode = (i == 0) ? llvm_ir::UnrollMode::kDefaultUnroll
                                 : llvm_ir::UnrollMode::kNoUnroll;
+
     if (bounds_index < dynamic_loop_bounds_->size()) {
       // Emit dynamic loop bounds for this dimension. Dynamic loop bounds
       // are read from ir function dynamic loop bounds argument.
@@ -69,14 +70,17 @@ ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock(absl::string_view loop_name,
 
       std::unique_ptr<llvm_ir::ForLoop> loop = loop_nest.AddLoop(
           /*suffix=*/absl::StrFormat("dim.%d", dimension), start_index,
-          end_index, unroll_mode);
+          end_index, unroll_mode, /*prevent_vectorization*/ false,
+          /* expression */ shape_.expressions(dimension));
       array_multi_index[dimension] = loop->GetIndVarValue();
     } else {
       // Emit static loop bounds for this dimension.
       std::unique_ptr<llvm_ir::ForLoop> loop = loop_nest.AddLoop(
           /*start_index=*/0,
           /*end_index=*/shape_.dimensions(dimension),
-          /*suffix=*/absl::StrFormat("dim.%d", dimension), unroll_mode);
+          /*suffix=*/absl::StrFormat("dim.%d", dimension), unroll_mode,
+          /*prevent_vectorization*/ false,
+          /* expression */ shape_.expressions(dimension));
       array_multi_index[dimension] = loop->GetIndVarValue();
     }
   }
diff --git a/third_party/xla/xla/service/cpu/runtime_batch_size.cc b/third_party/xla/xla/service/cpu/runtime_batch_size.cc
new file mode 100644
index 00000000000000..063cd5945b9a61
--- /dev/null
+++ b/third_party/xla/xla/service/cpu/runtime_batch_size.cc
@@ -0,0 +1,25 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/service/cpu/runtime_batch_size.h"
+
+#include "xla/executable_run_options.h"
+
+extern "C" int64_t __xla_cpu_runtime_GetBatchSize(
+    const void* run_options_ptr) {
+  if (run_options_ptr == nullptr) return 0;
+  return static_cast<const xla::ExecutableRunOptions*>(run_options_ptr)
+      ->batch_size();
+}
diff --git a/third_party/xla/xla/service/cpu/runtime_batch_size.h b/third_party/xla/xla/service/cpu/runtime_batch_size.h
new file mode 100644
index 00000000000000..fbd36294a4ddaf
--- /dev/null
+++ b/third_party/xla/xla/service/cpu/runtime_batch_size.h
@@ -0,0 +1,31 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_SERVICE_CPU_RUNTIME_BATCH_SIZE_H_
+#define XLA_SERVICE_CPU_RUNTIME_BATCH_SIZE_H_
+
+#include <cstdint>
+
+namespace xla::cpu::runtime {
+
+inline constexpr char kGetBatchSizeSymbolName[] =
+    "__xla_cpu_runtime_GetBatchSize";
+
+}  // namespace xla::cpu::runtime
+
+extern "C" int64_t __xla_cpu_runtime_GetBatchSize(
+    const void* run_options_ptr);
+
+#endif  // XLA_SERVICE_CPU_RUNTIME_BATCH_SIZE_H_
diff --git a/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc b/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc
index 87aca6c386751a..366655036cc038 100644
--- a/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc
+++ b/third_party/xla/xla/service/cpu/runtime_symbol_generator.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "llvm/Support/Error.h"
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 #include "xla/service/cpu/cpu_runtime.h"
+#include "xla/service/cpu/runtime_batch_size.h"
 #include "xla/service/cpu/runtime_conv2d.h"
 #include "xla/service/cpu/runtime_conv2d_acl.h"
 #include "xla/service/cpu/runtime_conv2d_mkl.h"
@@ -160,6 +161,7 @@ static bool RegisterKnownJITSymbols() {
   registry->Register("puts", reinterpret_cast<void*>(&puts), "Host");
 
   REGISTER_CPU_RUNTIME_SYMBOL(AcquireInfeedBufferForDequeue);
+  REGISTER_CPU_RUNTIME_SYMBOL(GetBatchSize);
   REGISTER_CPU_RUNTIME_SYMBOL(AcquireOutfeedBufferForPopulation);
   REGISTER_CPU_RUNTIME_SYMBOL(AllReduce);
   REGISTER_CPU_RUNTIME_SYMBOL(CollectivePermute);
@@ -201,6 +203,7 @@ static bool RegisterKnownJITSymbols() {
   REGISTER_CPU_RUNTIME_SYMBOL(EigenSingleThreadedMatMulU8);
   REGISTER_CPU_RUNTIME_SYMBOL(ParallelForkJoin);
   REGISTER_CPU_RUNTIME_SYMBOL(PrintfToStderr);
+  REGISTER_CPU_RUNTIME_SYMBOL(ReadCycleCounter);
   REGISTER_CPU_RUNTIME_SYMBOL(ReleaseInfeedBufferAfterDequeue);
   REGISTER_CPU_RUNTIME_SYMBOL(ReleaseOutfeedBufferAfterPopulation);
   REGISTER_CPU_RUNTIME_SYMBOL(StatusIsSuccess);
diff --git a/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc
index f0148de9932561..145beb3948b22d 100644
--- a/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/cpu_dyn_shape_test.cc
@@ -18,6 +18,8 @@ limitations under the License.
 #include <utility>
 
 #include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
@@ -32,7 +34,68 @@ namespace xla {
 namespace cpu {
 namespace {
 
-using CpuDynamicShapeTest = CpuCodegenTest;
+CpuAotCompilationOptions GetAotOptions() {
+  return CpuAotCompilationOptions{
+      /*triple=*/kTargetTripleForHost, /*cpu_name=*/kTargetCpuForHost,
+      /*features=*/"",
+      /*entry_point_name=*/"entry",
+      /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static};
+}
+
+Shape MakeDynamicDenseShape(absl::Span<const int64_t> dimensions,
+                            absl::Span<const int64_t> minor_to_major,
+                            int64_t dynamic_dimension) {
+  Shape shape =
+      ShapeUtil::MakeShapeWithDenseLayout(F32, dimensions, minor_to_major);
+  shape.set_dynamic_dimension(dynamic_dimension, true);
+  return shape;
+}
+
+class CpuDynamicShapeTest : public CpuCodegenTest {
+ protected:
+  std::unique_ptr<HloModule> CreateDynamicFastConcatModule(
+      absl::string_view name, bool add_consumer, bool three_operands = false) {
+    HloComputation::Builder builder(name);
+
+    Shape operand_shape = MakeDynamicDenseShape({64, 26, 4}, {2, 1, 0},
+                                                /*dynamic_dimension=*/0);
+    Shape concat_shape =
+        MakeDynamicDenseShape({64, 26, three_operands ? 12 : 8}, {2, 1, 0},
+                              /*dynamic_dimension=*/0);
+
+    HloInstruction* lhs = builder.AddInstruction(
+        HloInstruction::CreateParameter(0, operand_shape, "lhs"));
+    HloInstruction* rhs = builder.AddInstruction(
+        HloInstruction::CreateParameter(1, operand_shape, "rhs"));
+
+    HloInstruction* concat = nullptr;
+    if (three_operands) {
+      HloInstruction* third = builder.AddInstruction(
+          HloInstruction::CreateParameter(2, operand_shape, "third"));
+      concat = builder.AddInstruction(HloInstruction::CreateConcatenate(
+          concat_shape, {lhs, rhs, third}, /*dimension=*/2));
+    } else {
+      concat = builder.AddInstruction(HloInstruction::CreateConcatenate(
+          concat_shape, {lhs, rhs}, /*dimension=*/2));
+    }
+
+    if (add_consumer) {
+      int64_t addend_param_number = three_operands ? 3 : 2;
+      HloInstruction* addend =
+          builder.AddInstruction(HloInstruction::CreateParameter(
+              addend_param_number, concat_shape, "addend"));
+      builder.AddInstruction(HloInstruction::CreateBinary(
+          concat_shape, HloOpcode::kAdd, concat, addend));
+    }
+
+    auto hlo_module = CreateNewVerifiedModule();
+    hlo_module->AddEntryComputation(builder.Build());
+    auto& debug_options = hlo_module->mutable_config().mutable_debug_options();
+    debug_options.set_xla_cpu_use_thunk_runtime(false);
+    debug_options.add_xla_disable_hlo_passes("fusion");
+    return hlo_module;
+  }
+};
 
 TEST_F(CpuDynamicShapeTest, DynamicShapeR2) {
   HloComputation::Builder builder(TestName());
@@ -70,6 +133,56 @@ TEST_F(CpuDynamicShapeTest, DynamicShapeR2) {
                                 /*match_optimized_ir=*/false);
 }
 
+TEST_F(CpuDynamicShapeTest, DynamicFastConcatenateCompilesWithoutFusion) {
+  auto hlo_module =
+      CreateDynamicFastConcatModule(TestName(), /*add_consumer=*/false);
+
+  std::string filecheck_pattern = R"(
+; CHECK: %[[dyn_dim_size:.*]] = load i32, ptr
+; CHECK: %[[i64_dyn_dim_size:.*]] = sext i32 %[[dyn_dim_size:.*]] to i64
+; CHECK: br label %concatenate.loop_header.concat.1
+; CHECK: call void @llvm.memcpy
+  )";
+
+  CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), GetAotOptions(),
+                                filecheck_pattern,
+                                /*match_optimized_ir=*/false);
+}
+
+TEST_F(CpuDynamicShapeTest,
+       DynamicFastConcatenateWithAddCompilesWithoutFusion) {
+  auto hlo_module =
+      CreateDynamicFastConcatModule(TestName(), /*add_consumer=*/true);
+
+  std::string filecheck_pattern = R"(
+; CHECK: %[[dyn_dim_size:.*]] = load i32, ptr
+; CHECK: %[[i64_dyn_dim_size:.*]] = sext i32 %[[dyn_dim_size:.*]] to i64
+; CHECK: br label %concatenate.loop_header.concat.1
+; CHECK: call void @llvm.memcpy
+  )";
+
+  CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), GetAotOptions(),
+                                filecheck_pattern,
+                                /*match_optimized_ir=*/false);
+}
+
+TEST_F(CpuDynamicShapeTest,
+       DynamicFastConcatenateWithThreeOperandsCompilesWithoutFusion) {
+  auto hlo_module = CreateDynamicFastConcatModule(
+      TestName(), /*add_consumer=*/true, /*three_operands=*/true);
+
+  std::string filecheck_pattern = R"(
+; CHECK: %[[dyn_dim_size:.*]] = load i32, ptr
+; CHECK: %[[i64_dyn_dim_size:.*]] = sext i32 %[[dyn_dim_size:.*]] to i64
+; CHECK: br label %concatenate.loop_header.concat.1
+; CHECK: call void @llvm.memcpy
+  )";
+
+  CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), GetAotOptions(),
+                                filecheck_pattern,
+                                /*match_optimized_ir=*/false);
+}
+
 }  // namespace
 }  // namespace cpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.cc b/third_party/xla/xla/service/cpu/thunk_emitter.cc
index 9079ed8a16d1a8..72fe6b95e25825 100644
--- a/third_party/xla/xla/service/cpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/cpu/thunk_emitter.cc
@@ -1086,6 +1086,8 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCustomCallThunk(
     return EmitTopKThunk(custom_call);
   } else if (custom_call_target == "SliceToDynamic") {
     return EmitSliceToDynamicThunk(instruction);
+  } else if (custom_call_target == "GetExpressionValue") {
+    return EmitGetExpressionValueThunk(instruction);
   }
 
   // Check the API version.
@@ -1128,6 +1130,21 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitSliceToDynamicThunk(
       /*min_alignment=*/cpu_function_runtime::MinAlign());
 }
 
+absl::StatusOr<ThunkSequence> ThunkEmitter::EmitGetExpressionValueThunk(
+    const HloInstruction* instruction) {
+  VLOG(2) << "Handling GetExpressionValue for instruction: "
+          << instruction->ToString();
+  const HloCustomCallInstruction* custom_call =
+      Cast<HloCustomCallInstruction>(instruction);
+  TF_ASSIGN_OR_RETURN(
+      auto kernel, ir_emitter_.EmitGetExpressionValueHostKernel(custom_call));
+  TF_ASSIGN_OR_RETURN(auto result_buffer,
+                      GetHostKernelAllocationSlices(instruction));
+  return MakeKernelThunkSequence(
+      instruction, result_buffer, kernel,
+      /*min_alignment=*/cpu_function_runtime::MinAlign());
+}
+
 absl::StatusOr<ThunkSequence> ThunkEmitter::EmitSliceThunk(
     const HloInstruction* instruction) {
   // TODO(ezhulenev): Consider implementing slice operations as separate
diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.h b/third_party/xla/xla/service/cpu/thunk_emitter.h
index 9cdc8ae3981680..40aad964fb8d24 100644
--- a/third_party/xla/xla/service/cpu/thunk_emitter.h
+++ b/third_party/xla/xla/service/cpu/thunk_emitter.h
@@ -183,6 +183,9 @@ class ThunkEmitter {
   absl::StatusOr<ThunkSequence> EmitSliceToDynamicThunk(
       const HloInstruction* instruction);
 
+  absl::StatusOr<ThunkSequence> EmitGetExpressionValueThunk(
+      const HloInstruction* instruction);
+
   absl::StatusOr<ThunkSequence> EmitTopKThunk(
       const HloCustomCallInstruction* custom_call);
 
diff --git a/third_party/xla/xla/service/dynamic_constant_rewriter.cc b/third_party/xla/xla/service/dynamic_constant_rewriter.cc
new file mode 100644
index 00000000000000..3c734006456b79
--- /dev/null
+++ b/third_party/xla/xla/service/dynamic_constant_rewriter.cc
@@ -0,0 +1,165 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/service/dynamic_constant_rewriter.h"
+
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/literal_util.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/status_macros.h"
+#include "tsl/platform/errors.h"
+
+namespace xla {
+namespace {
+
+absl::StatusOr<HloInstruction*> BuildDynamicConstantReplacement(
+    HloInstruction* constant_instr) {
+  TF_RET_CHECK(constant_instr->opcode() == HloOpcode::kConstant);
+  TF_RET_CHECK(constant_instr->has_contents());
+
+  const Shape& shape = constant_instr->shape();
+  TF_RET_CHECK(shape.IsArray());
+  TF_RET_CHECK(shape.element_type() == S32 || shape.element_type() == S64)
+      << "Only s32/s64 marked constants are supported";
+  TF_RET_CHECK(shape.dimensions_size() <= 1)
+      << "Only scalar and rank-1 marked constants are supported";
+
+  int64_t dynamic_index = 0;
+  DExpr expr;
+  for (int64_t i = 0; i < constant_instr->contents().size(); ++i) {
+    DExpr candidate = DExprFromProto(constant_instr->contents()[i]);
+    if (candidate && candidate->is_dynamic()) {
+      dynamic_index = i;
+      expr = std::move(candidate);
+      break;
+    }
+  }
+  TF_RET_CHECK(expr) << "Marked dynamic constant is missing dynamic contents";
+
+  int64_t carrier_bound;
+  if (shape.dimensions_size() == 0) {
+    dynamic_index = 0;
+    carrier_bound = shape.element_type() == S32
+                        ? constant_instr->literal().GetFirstElement<int32_t>()
+                        : constant_instr->literal().GetFirstElement<int64_t>();
+  } else {
+    TF_RET_CHECK(dynamic_index >= 0 && dynamic_index < shape.dimensions(0))
+        << "dynamic content index=" << dynamic_index
+        << " out of bounds for shape " << shape.ToString();
+    carrier_bound = shape.element_type() == S32
+                        ? constant_instr->literal().data<int32_t>()[dynamic_index]
+                        : constant_instr->literal().data<int64_t>()[dynamic_index];
+  }
+
+  TF_RET_CHECK(carrier_bound <= std::numeric_limits<int32_t>::max() &&
+               carrier_bound >= std::numeric_limits<int32_t>::min())
+      << "GetExpressionValue carriers must fit in s32, got " << carrier_bound;
+
+  HloComputation* computation = constant_instr->parent();
+  HloInstruction* carrier = computation->AddInstruction(
+      HloInstruction::CreateConstant(
+          LiteralUtil::CreateR1<int32_t>(
+              {static_cast<int32_t>(carrier_bound)})));
+  ExpressionProto expr_proto;
+  expr.to_proto(&expr_proto);
+
+  HloInstruction* runtime_value = computation->AddInstruction(
+      HloInstruction::CreateCustomCall(
+          ShapeUtil::MakeShape(S32, {}), {carrier}, "GetExpressionValue"));
+  runtime_value->set_contents({std::move(expr_proto)});
+  if (shape.element_type() == S64) {
+    runtime_value = computation->AddInstruction(HloInstruction::CreateConvert(
+        ShapeUtil::MakeShape(S64, {}), runtime_value));
+  }
+
+  if (shape.dimensions_size() == 0) {
+    return runtime_value;
+  }
+
+  HloInstruction* base_constant =
+      computation->AddInstruction(constant_instr->Clone());
+  base_constant->set_contents({});
+  Shape update_shape = ShapeUtil::MakeShape(shape.element_type(), {1});
+  HloInstruction* update = computation->AddInstruction(
+      HloInstruction::CreateReshape(update_shape, runtime_value));
+  HloInstruction* start_index = computation->AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateR0<int32_t>(
+          static_cast<int32_t>(dynamic_index))));
+  return computation->AddInstruction(HloInstruction::CreateDynamicUpdateSlice(
+      shape, base_constant, update, {start_index}));
+}
+
+bool IsMarkedDynamicConstant(const HloInstruction* instr) {
+  return instr->opcode() == HloOpcode::kConstant && instr->has_contents();
+}
+
+}  // namespace
+
+absl::StatusOr<bool> DynamicConstantRewriter::Run(
+    HloModule* module,
+    const absl::flat_hash_set<absl::string_view>& execution_threads) {
+  bool changed = false;
+  for (HloComputation* computation :
+       module->MakeNonfusionComputations(execution_threads)) {
+    std::vector<HloInstruction*> marked_constants;
+    for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) {
+      if (instruction->opcode() == HloOpcode::kConstant) {
+        const bool is_marked = IsMarkedDynamicConstant(instruction);
+        VLOG(1) << "DynamicConstantRewriter sees constant "
+                << instruction->name() << " shape="
+                << instruction->shape().ToString()
+                << " literal=" << instruction->literal().ToString()
+                << " marked=" << is_marked;
+        if (is_marked) {
+          VLOG(1) << "  contents_size=" << instruction->contents().size();
+          marked_constants.push_back(instruction);
+        }
+      }
+    }
+
+    for (HloInstruction* constant_instr : marked_constants) {
+      if (constant_instr->IsDead()) {
+        continue;
+      }
+      TF_ASSIGN_OR_RETURN(HloInstruction * replacement,
+                          BuildDynamicConstantReplacement(constant_instr));
+      VLOG(1) << "Rewriting marked constant " << constant_instr->name()
+              << " into " << replacement->name() << " ("
+              << HloOpcodeString(replacement->opcode()) << ")";
+      TF_RETURN_IF_ERROR(
+          computation->ReplaceInstruction(constant_instr, replacement));
+      replacement->erase_frontend_attribute("dynamic_constant_index");
+      replacement->erase_frontend_attribute("dynamic_constant_expr");
+      changed = true;
+    }
+  }
+  return changed;
+}
+
+}  // namespace xla
diff --git a/third_party/xla/xla/service/dynamic_constant_rewriter.h b/third_party/xla/xla/service/dynamic_constant_rewriter.h
new file mode 100644
index 00000000000000..e1bd5c57cadf5c
--- /dev/null
+++ b/third_party/xla/xla/service/dynamic_constant_rewriter.h
@@ -0,0 +1,41 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_SERVICE_DYNAMIC_CONSTANT_REWRITER_H_
+#define XLA_SERVICE_DYNAMIC_CONSTANT_REWRITER_H_
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
+
+namespace xla {
+
+class DynamicConstantRewriter : public HloModulePass {
+ public:
+  absl::string_view name() const override {
+    return "dynamic_constant_rewriter";
+  }
+
+  using HloPassInterface::Run;
+  absl::StatusOr<bool> Run(
+      HloModule* module,
+      const absl::flat_hash_set<absl::string_view>& execution_threads) override;
+};
+
+}  // namespace xla
+
+#endif  // XLA_SERVICE_DYNAMIC_CONSTANT_REWRITER_H_
diff --git a/third_party/xla/xla/service/elemental_ir_emitter.cc b/third_party/xla/xla/service/elemental_ir_emitter.cc
index 89391a3bdd1dd5..3fe608125e8c3e 100644
--- a/third_party/xla/xla/service/elemental_ir_emitter.cc
+++ b/third_party/xla/xla/service/elemental_ir_emitter.cc
@@ -3278,51 +3278,57 @@ absl::StatusOr<llvm::Value*> ElementalIrEmitter::EmitElementalConcatenate(
   }
 
   // We use bisection to select the input operand.
-  int64_t current_offset = 0;
+  int64_t coffset = 0;
+  llvm::Value* current_offset = source_index.GetConstantWithIndexType(0);
 
   // Offset for every operand.
-  std::vector<std::pair<int64_t, const HloInstruction*>> cases;
+  std::vector<std::pair<llvm::Value*, const HloInstruction*>> cases;
 
   cases.reserve(hlo->operand_count());
   for (const HloInstruction* operand : hlo->operands()) {
     cases.emplace_back(current_offset, operand);
-    current_offset += operand->shape().dimensions(concat_dim);
+    llvm::Value* cdim = source_index.GetConstantWithIndexType(
+        operand->shape().dimensions(concat_dim));
+    const auto& concat_expr = operand->shape().expressions(concat_dim);
+    if (concat_expr->is_dynamic()) {
+      cdim = llvm_ir::EmitExpression(b_, concat_expr);
+    }
+    current_offset = b_->CreateAdd(current_offset, cdim, "current_offset");
+    coffset += operand->shape().dimensions(concat_dim);
   }
-  CHECK_EQ(current_offset, hlo->shape().dimensions(concat_dim));
+  CHECK_EQ(coffset, hlo->shape().dimensions(concat_dim));
 
   std::function<llvm::BasicBlock*(
-      absl::Span<const std::pair<int64_t, const HloInstruction*>> operands)>
+      absl::Span<const std::pair<llvm::Value*, const HloInstruction*>> operands)>
       emit_tree =
-          [&](absl::Span<const std::pair<int64_t, const HloInstruction*>>
+          [&](absl::Span<const std::pair<llvm::Value*, const HloInstruction*>>
                   operands) {
             llvm::IRBuilder<>::InsertPointGuard guard(*b_);
             size_t mid = operands.size() / 2;
-            const std::pair<int64_t, const HloInstruction*>& pivot =
+            const std::pair<llvm::Value*, const HloInstruction*>& pivot =
                 operands[mid];
             llvm::BasicBlock* block = llvm_ir::CreateBasicBlock(
                 exit_block,
-                absl::StrCat("concatenate.pivot.", pivot.first, "."), b_);
+                absl::StrCat("concatenate.pivot."), b_);
             b_->SetInsertPoint(block);
 
             // If there's only one element we're done. The range is contiguous
             // so we can just jump to the block for it.
             if (operands.size() == 1) {
-              const std::pair<int64_t, const HloInstruction*>& operand =
+              const std::pair<llvm::Value*, const HloInstruction*>& operand =
                   operands.back();
               int64_t operand_id = to_unique_operand_id[operand.second];
 
               source_index_phis[operand_id]->addIncoming(
-                  source_index.GetConstantWithIndexType(operand.first),
+                  operand.first,
                   b_->GetInsertBlock());
               b_->CreateBr(emit_operand_blocks[operand_id]);
               return block;
             }
 
             // Take the middle element and recurse.
-            llvm::Constant* pivot_const = llvm::ConstantInt::get(
-                source_index[concat_dim]->getType(), pivot.first);
             llvm::Value* comp =
-                b_->CreateICmpULT(source_index[concat_dim], pivot_const);
+                b_->CreateICmpULT(source_index[concat_dim], pivot.first);
 
             llvm::BasicBlock* left_block = emit_tree(operands.subspan(0, mid));
             llvm::BasicBlock* right_block = emit_tree(operands.subspan(mid));
@@ -3616,11 +3622,16 @@ absl::StatusOr<llvm::Value*> ElementalIrEmitter::EmitElementalPad(
             "in_bounds");
     multi_index[i] =
         SDiv(multi_index[i], index_typed_const(pad_dim.interior_padding() + 1));
-    in_bounds =
-        And(in_bounds,
-            ICmpSLT(multi_index[i],
-                    index_typed_const(hlo->operand(0)->shape().dimensions(i))),
-            "in_bounds");
+
+    int64_t shape_dim = hlo->operand(0)->shape().dimensions(i);
+    llvm::Value* bound = index_typed_const(shape_dim);
+
+    const auto& operand_expr = hlo->operand(0)->shape().expressions(i);
+    if (operand_expr->is_dynamic()) {
+      bound = llvm_ir::EmitExpression(b_, operand_expr);
+    }
+
+    in_bounds = And(in_bounds, ICmpSLT(multi_index[i], bound), "in_bounds");
   }
 
   // if (in_bounds) {
@@ -3687,9 +3698,20 @@ absl::StatusOr<llvm::Value*> ElementalIrEmitter::EmitElementalDot(
     return llvm::ConstantInt::get(index_type, c);
   };
 
+  llvm::Value* contracted_bound = index_typed_const(contracted_dim_size);
+
+  if (!hlo->operand(0)
+           ->shape()
+           .expressions(lhs_contracting_dim)
+           ->is_constant()) {
+    llvm::Value* expr_value = llvm_ir::EmitExpression(
+        b_, hlo->operand(0)->shape().expressions(lhs_contracting_dim));
+    contracted_bound = expr_value;
+  }
+
   std::unique_ptr<llvm_ir::ForLoop> inner_loop = llvm_ir::ForLoop::EmitForLoop(
-      IrName(hlo, "inner"), index_typed_const(0),
-      index_typed_const(contracted_dim_size), index_typed_const(1), b_);
+      IrName(hlo, "inner"), index_typed_const(0), contracted_bound,
+      index_typed_const(1), b_);
 
   SetToFirstInsertPoint(inner_loop->GetPreheaderBasicBlock(), b_);
   PrimitiveType primitive_type = hlo->shape().element_type();
@@ -3883,9 +3905,17 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator(
         const HloInstruction* operand = hlo->operand(0);
         std::vector<llvm::Value*> source_multi_index = target_index.multidim();
         for (int64_t dim : hlo->dimensions()) {
-          source_multi_index[dim] = Sub(target_index.GetConstantWithIndexType(
-                                            hlo->shape().dimensions(dim) - 1),
-                                        target_index[dim]);
+          const auto& dim_expr = hlo->shape().expressions(dim);
+          if (dim_expr->is_dynamic()) {
+            llvm::Value* one = target_index.GetConstantWithIndexType(1);
+            llvm::Value* expr_value = llvm_ir::EmitExpression(b_, dim_expr);
+            source_multi_index[dim] =
+                Sub(Sub(expr_value, one), target_index[dim]);
+          } else {
+            source_multi_index[dim] = Sub(target_index.GetConstantWithIndexType(
+                                              hlo->shape().dimensions(dim) - 1),
+                                          target_index[dim]);
+          }
         }
         llvm_ir::IrArray::Index source_index(
             source_multi_index, operand->shape(), target_index.GetType());
@@ -3967,11 +3997,12 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator(
     case HloOpcode::kSlice:
       return [this, hlo, &operand_to_generator](
                  const IrArray::Index& index) -> absl::StatusOr<llvm::Value*> {
-        IrArray::Index sliced_index = index.SourceIndexOfSlice(
-            /*operand_shape=*/hlo->operand(0)->shape(),
-            /*starts=*/hlo->slice_starts(),
-            /*strides=*/hlo->slice_strides(), /*builder=*/b_);
-        return operand_to_generator.at(hlo->operand(0))(sliced_index);
+        const auto* slice = Cast<HloSliceInstruction>(hlo);
+        return operand_to_generator.at(hlo->operand(0))(
+            index.SourceIndexOfSlice(hlo->operand(0)->shape(),
+                                     hlo->slice_starts(),
+                                     slice->slice_start_exprs(),
+                                     hlo->slice_strides(), b_));
       };
     case HloOpcode::kDynamicSlice:
       return [this, hlo, &operand_to_generator](
@@ -4236,11 +4267,21 @@ absl::StatusOr<llvm::Value*> ElementalIrEmitter::EmitElementalReduceWindow(
     // comparison is equivalent to the unsigned comparison
     // input_multi_index[i] < bound, as a negative value wraps to a large
     // positive value.
+
+    int64_t dim_bound = reduce_window->inputs()[0]->shape().dimensions(i);
+    llvm::Value* shape_bound = index_typed_const(dim_bound);
+
+    const auto& window_expr =
+        reduce_window->inputs()[0]->shape().expressions(i);
+    if (window_expr->is_dynamic()) {
+      llvm::Value* expr_value = llvm_ir::EmitExpression(b_, window_expr);
+      shape_bound = expr_value;
+    }
+
     in_bounds =
         And(in_bounds,
             ICmpULT(input_multi_index[i],
-                    index_typed_const(
-                        reduce_window->inputs()[0]->shape().dimensions(i))));
+                    shape_bound));
   }
 
   llvm_ir::LlvmIfData if_data =
diff --git a/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc b/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc
index 5054a440778105..484a1e75ac0aeb 100644
--- a/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc
+++ b/third_party/xla/xla/service/gpu/transforms/gemm_rewriter.cc
@@ -1316,6 +1316,7 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
         x = instr->AddInstruction(op.first->CloneWithNewOperands(
             ShapeUtil::MakeShapeWithDenseLayout(
                 x->shape().element_type(), op.first->shape().dimensions(),
+                op.first->shape().expressions(),
                 op.first->shape().layout().minor_to_major()),
             operands));
       }
@@ -1378,6 +1379,7 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
         instr->AddInstruction(HloInstruction::CreateCustomCall(
             ShapeUtil::MakeShapeWithDenseLayout(
                 instr->shape().element_type(), new_output_shape.dimensions(),
+                new_output_shape.expressions(),
                 instr->shape().layout().minor_to_major()),
             operands_list, kCublasLtMatmulF8CallTarget));
     TF_RETURN_IF_ERROR(new_custom_call->set_backend_config(gpu_backend_config));
diff --git a/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc b/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc
index ce454624144803..4585af7203946b 100644
--- a/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc
+++ b/third_party/xla/xla/service/gpu/transforms/windowed_einsum_handler.cc
@@ -183,11 +183,13 @@ absl::StatusOr<HloInstruction*> ShiftDequantizationF8(
     for (HloInstruction* unary : unaries[k]) {
       Shape new_shape = ShapeUtil::MakeShapeWithDenseLayout(
           operands[k]->shape().element_type(), unary->shape().dimensions(),
+          unary->shape().expressions(),
           unary->shape().layout().minor_to_major());
 
       operands[k] = unary->AddInstruction(unary->CloneWithNewOperands(
           ShapeUtil::MakeShapeWithDenseLayout(
               operands[k]->shape().element_type(), unary->shape().dimensions(),
+              unary->shape().expressions(),
               unary->shape().layout().minor_to_major()),
           {operands[k]}));
     }
diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto
index 713db5886c6f4e..de4c27b9367843 100644
--- a/third_party/xla/xla/service/hlo.proto
+++ b/third_party/xla/xla/service/hlo.proto
@@ -113,7 +113,7 @@ enum CustomCallApiVersion {
 }
 
 // Serialization of HloInstruction.
-// Next ID: 92
+// Next ID: 93
 message HloInstructionProto {
   reserved 10;
   reserved "parameter_name";
@@ -174,6 +174,8 @@ message HloInstructionProto {
     int64 start = 1;
     int64 limit = 2;
     int64 stride = 3;
+    xla.ExpressionProto start_expr = 4;
+    xla.ExpressionProto limit_expr = 5;
   }
   repeated SliceDimensions slice_dimensions = 17;
 
@@ -328,6 +330,9 @@ message HloInstructionProto {
   // Frontend attributes to pass to the XLA backend.
   xla.FrontendAttributes frontend_attributes = 68;
 
+  // Structured symbolic contents attached to this instruction.
+  repeated xla.ExpressionProto contents = 92;
+
   // Specifies if all elements updated are guaranteed to be unique by
   // the caller.
   bool unique_indices = 69;
diff --git a/third_party/xla/xla/service/hlo_creation_utils.cc b/third_party/xla/xla/service/hlo_creation_utils.cc
index b815b1ffc79627..15c2c8417c9f28 100644
--- a/third_party/xla/xla/service/hlo_creation_utils.cc
+++ b/third_party/xla/xla/service/hlo_creation_utils.cc
@@ -184,7 +184,8 @@ absl::StatusOr<HloInstruction*> MakeDynamicSliceHlo(
   TF_ASSIGN_OR_RETURN(
       Shape dynamic_slice_shape,
       ShapeInference::InferDynamicSliceShape(
-          operand->shape(), scalar_start_indices_shapes, slice_sizes));
+          operand->shape(), scalar_start_indices_shapes, slice_sizes,
+          operand->shape().expressions()));
   return computation->AddInstruction(
       HloInstruction::CreateDynamicSlice(dynamic_slice_shape, operand,
                                          start_indices, slice_sizes),
@@ -655,9 +656,12 @@ absl::StatusOr<HloInstruction*> CollapseFirstNDims(HloInstruction* operand,
   CHECK_GE(operand_shape.dimensions_size(), n);
   int64_t new_shape_leading_bound = 1;
   bool new_shape_leading_is_dynamic = false;
+  DExpr new_shape_leading_expression = DExpr::Const(1);
   for (int64_t i = 0; i < n; i++) {
     new_shape_leading_bound *= operand_shape.dimensions(i);
     new_shape_leading_is_dynamic |= operand_shape.is_dynamic_dimension(i);
+    new_shape_leading_expression =
+        new_shape_leading_expression * operand_shape.expressions(i);
   }
 
   std::vector<int64_t> new_shape_dims;
@@ -675,8 +679,16 @@ absl::StatusOr<HloInstruction*> CollapseFirstNDims(HloInstruction* operand,
             operand_shape.dynamic_dimensions().end(),
             std::back_inserter(new_shape_dynamic_dims));
 
-  Shape output_shape = ShapeUtil::MakeShape(
-      operand_shape.element_type(), new_shape_dims, new_shape_dynamic_dims);
+  std::vector<DExpr> new_shape_expressions;
+  new_shape_expressions.reserve(operand_shape.dimensions_size() - n + 1);
+  new_shape_expressions.push_back(new_shape_leading_expression);
+  auto exprs = operand_shape.expressions();
+  std::copy(exprs.begin() + n, exprs.end(),
+            std::back_inserter(new_shape_expressions));
+
+  Shape output_shape =
+      ShapeUtil::MakeShape(operand_shape.element_type(), new_shape_dims,
+                           new_shape_dynamic_dims, new_shape_expressions);
 
   return MakeReshapeHlo(output_shape, operand);
 }
diff --git a/third_party/xla/xla/service/hlo_cse.cc b/third_party/xla/xla/service/hlo_cse.cc
index 570ffa03de0edf..4b3ad32340cd44 100644
--- a/third_party/xla/xla/service/hlo_cse.cc
+++ b/third_party/xla/xla/service/hlo_cse.cc
@@ -37,6 +37,26 @@ namespace xla {
 
 namespace {
 
+bool HasDynamicConstantContents(const HloInstruction* instruction) {
+  if (!instruction->has_contents()) {
+    return false;
+  }
+  for (const auto& content : instruction->contents()) {
+    if (content.node_type_case() != ExpressionProto::kConstantValue &&
+        content.node_type_case() != ExpressionProto::NODE_TYPE_NOT_SET) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool HasDynamicConstantFrontendAttributes(const HloInstruction* instruction) {
+  const auto& attrs = instruction->frontend_attributes().map();
+  return HasDynamicConstantContents(instruction) ||
+         attrs.contains("dynamic_constant_index") ||
+         attrs.contains("dynamic_constant_expr");
+}
+
 template <bool kIsLayoutSensitive>
 struct ConstantKey {
   template <typename H>
@@ -92,6 +112,13 @@ absl::StatusOr<bool> CombineConstants(HloComputation* computation,
       continue;
     }
 
+    // Dynamic constant frontend attrs carry semantic metadata for a later
+    // rewrite pass, so these constants are not interchangeable with otherwise
+    // identical plain constants.
+    if (HasDynamicConstantFrontendAttributes(instruction)) {
+      continue;
+    }
+
     HloInstruction* match = nullptr;
     if (auto* constant_inst = DynCast<HloConstantInstruction>(instruction)) {
       auto insert_result = constants.insert(ConstantKey<kIsLayoutSensitive>{
@@ -310,6 +337,12 @@ absl::StatusOr<bool> HloCSE::RunOnComputation(HloComputation* computation) {
       continue;
     }
 
+    // These frontend attrs are semantic, not just decorative, so do not CSE
+    // across them.
+    if (HasDynamicConstantFrontendAttributes(instruction)) {
+      continue;
+    }
+
     auto pair = representatives.insert(CseKey{instruction});
     if (!pair.second) {
       HloInstruction* equivalent_instruction = pair.first->hlo;
diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index b7bea1e55704c5..001a873118ab7f 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -1401,6 +1401,7 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOperandLayoutFromOutputLayout(
     const Shape& output_shape = instruction->shape();
     Shape output_shape_with_layout = ShapeUtil::MakeShapeWithDenseLayout(
         output_shape.element_type(), output_shape.dimensions(),
+        output_shape.expressions(),
         LayoutUtil::MinorToMajor(output_layout));
     Shape operand_shape = operand->shape();
     *operand_shape.mutable_layout() =
@@ -1539,6 +1540,7 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOutputLayoutFromOperandLayout(
     }
     Shape operand_shape_with_layout = ShapeUtil::MakeShapeWithDenseLayout(
         operand->shape().element_type(), operand->shape().dimensions(),
+        operand->shape().expressions(),
         LayoutUtil::MinorToMajor(operand_layout));
     Shape output_shape = user->shape();
     *output_shape.mutable_layout() =
@@ -2571,7 +2573,7 @@ absl::Status LayoutAssignment::PropagateComputationLayouts(
     *result_layout = computed_computation_layout.result_layout();
   } else {
     TF_RET_CHECK(
-        Shape::Equal().IgnoreDynamicDimension().MinorToMajorOnlyInLayout()(
+        Shape::Equal().IgnoreDynamicDimension().MinorToMajorOnlyInLayout().IgnoreBatch()(
             computed_computation_layout.result_layout().shape(),
             result_layout->shape()));
   }
diff --git a/third_party/xla/xla/service/llvm_ir/BUILD b/third_party/xla/xla/service/llvm_ir/BUILD
index 599060b88eee81..0f3cd53a985654 100644
--- a/third_party/xla/xla/service/llvm_ir/BUILD
+++ b/third_party/xla/xla/service/llvm_ir/BUILD
@@ -79,9 +79,10 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
+        "//xla/service/cpu:cpu_options",
+        "//xla/service/cpu:runtime_batch_size_api",
         "//xla/service:dump",
         "//xla/service:hlo_module_config",
-        "//xla/service/cpu:cpu_options",
         "//xla/tsl/platform:byte_order",
         "//xla/tsl/platform:logging",
         "@com_google_absl//absl/algorithm:container",
@@ -355,3 +356,18 @@ xla_cc_test(
         "@llvm-project//llvm:ir_headers",
     ],
 )
+
+xla_cc_test(
+    name = "llvm_util_test",
+    srcs = ["llvm_util_test.cc"],
+    deps = [
+        ":llvm_util",
+        "//xla/service/cpu:runtime_batch_size_api",
+        "//xla:shape_util",
+        "//xla/tests:xla_internal_test_main",
+        "@com_google_googletest//:gtest",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:ir_headers",
+    ],
+)
diff --git a/third_party/xla/xla/service/llvm_ir/ir_array.cc b/third_party/xla/xla/service/llvm_ir/ir_array.cc
index 139d537c88778a..8288cba9530899 100644
--- a/third_party/xla/xla/service/llvm_ir/ir_array.cc
+++ b/third_party/xla/xla/service/llvm_ir/ir_array.cc
@@ -281,8 +281,11 @@ IrArray::Index IrArray::Index::SourceIndexOfReshape(
       // linear index by each dimension size.
       for (int64_t i = common_factors[k + 1].first - 1;
            i >= common_factors[k].first; --i) {
+        const auto& input_expr = input_shape.expressions(i);
+        bool is_dynamic = input_expr->is_dynamic();
         llvm::Value* divisor =
-            GetConstantWithIndexType(input_shape.dimensions(i));
+            is_dynamic ? llvm_ir::EmitExpression(builder, input_expr)
+                       : GetConstantWithIndexType(input_shape.dimensions(i));
         if (input_shape.dimensions(i) == 1) {
           source_multidim_index[i] = GetConstantWithIndexType(0);
         } else if (i == common_factors[k].first) {
@@ -307,17 +310,27 @@ IrArray::Index IrArray::Index::SourceIndexOfReshape(
 
 IrArray::Index IrArray::Index::SourceIndexOfSlice(
     const Shape& operand_shape, absl::Span<const int64_t> starts,
-    absl::Span<const int64_t> strides, llvm::IRBuilderBase* builder) const {
+    absl::Span<const DExpr> start_exprs, absl::Span<const int64_t> strides,
+    llvm::IRBuilderBase* builder) const {
+  static const DExpr kMissingSliceExpr =
+      DExpr::Unknown(kMissingExpressionSentinel);
   std::vector<llvm::Value*> source_multi_index(multidim_.size());
   for (int i = 0; i < multidim_.size(); ++i) {
+    llvm::Value* start_value = GetConstantWithIndexType(starts[i]);
+    const DExpr& start_expr =
+        i < start_exprs.size() ? start_exprs[i] : kMissingSliceExpr;
+    if (start_expr && start_expr->is_dynamic()) {
+      start_value = builder->CreateIntCast(
+          llvm_ir::EmitExpression(builder, start_expr), index_type_,
+          /*isSigned=*/true);
+    }
     int64_t stride = strides[i];
     if (stride != 1) {
       source_multi_index[i] = builder->CreateAdd(
           builder->CreateMul(multidim_[i], GetConstantWithIndexType(stride)),
-          GetConstantWithIndexType(starts[i]));
+          start_value);
     } else {
-      source_multi_index[i] =
-          builder->CreateAdd(multidim_[i], GetConstantWithIndexType(starts[i]));
+      source_multi_index[i] = builder->CreateAdd(multidim_[i], start_value);
     }
   }
   return Index(source_multi_index, operand_shape, index_type_);
@@ -559,8 +572,33 @@ llvm::Value* IrArray::EmitArrayElementAddress(const IrArray::Index& index,
     int64_t dimension = LayoutUtil::Major(shape_.layout(), i);
     gep_indices.push_back(actual_index[dimension]);
   }
-  return b->CreateInBoundsGEP(pointee_type_, base_ptr_, gep_indices,
-                              llvm_ir::AsStringRef(name));
+
+  // Do not make a dynamic "GEP" if only the first dimension is dynamic since
+  // it's always indiced with 0 (i.e. the dynamic dimension has no impact on the
+  // address computation).
+  std::vector<int64_t> gep_dims;
+  std::vector<DExpr> gep_expressions;
+  gep_dims.reserve(shape_.dimensions().size());
+  gep_expressions.reserve(shape_.dimensions().size());
+  for (int64_t i = 0; i < shape_.dimensions().size(); ++i) {
+    int64_t dimension = LayoutUtil::Major(shape_.layout(), i);
+    gep_dims.push_back(shape_.dimensions(dimension));
+    gep_expressions.push_back(shape_.expressions(dimension));
+  }
+  bool dynamic_first_dim =
+      gep_expressions[0]->is_dynamic() &&
+      std::all_of(gep_expressions.begin() + 1, gep_expressions.end(),
+                  [](const DExpr& e) { return e->is_constant(); });
+  if (!dynamic_first_dim && shape_.has_dynamic_expr()) {
+    llvm::Type* element_type =
+        PrimitiveTypeToIrType(shape_.element_type(), b->getContext());
+    return llvm_ir::createDynamicGEP(
+        b, base_ptr_, gep_indices, gep_dims, gep_expressions,
+        element_type, llvm_ir::AsStringRef(name));
+  } else {
+    return b->CreateInBoundsGEP(pointee_type_, base_ptr_, gep_indices,
+                                llvm_ir::AsStringRef(name));
+  }
 }
 
 llvm::Value* IrArray::EmitLinearArrayElementAddress(
diff --git a/third_party/xla/xla/service/llvm_ir/ir_array.h b/third_party/xla/xla/service/llvm_ir/ir_array.h
index f5b2b7c4fbd792..eca880b39d29f1 100644
--- a/third_party/xla/xla/service/llvm_ir/ir_array.h
+++ b/third_party/xla/xla/service/llvm_ir/ir_array.h
@@ -158,6 +158,7 @@ class IrArray {
     // `operand_shape`.
     Index SourceIndexOfSlice(const Shape& operand_shape,
                              absl::Span<const int64_t> starts,
+                             absl::Span<const DExpr> start_exprs,
                              absl::Span<const int64_t> strides,
                              llvm::IRBuilderBase* builder) const;
 
diff --git a/third_party/xla/xla/service/llvm_ir/llvm_loop.cc b/third_party/xla/xla/service/llvm_ir/llvm_loop.cc
index 43ee77ef0e85be..c60fc1cc7fa8ba 100644
--- a/third_party/xla/xla/service/llvm_ir/llvm_loop.cc
+++ b/third_party/xla/xla/service/llvm_ir/llvm_loop.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/shape.h"
 #include "xla/tsl/platform/logging.h"
+#include "llvm/include/llvm/Support/Debug.h"
 
 namespace xla {
 namespace llvm_ir {
@@ -185,24 +186,40 @@ llvm::BasicBlock* ForLoop::CreateLoopBB(absl::string_view name,
   return CreateBasicBlock(insert_before_bb_, GetQualifiedName(name), b);
 }
 
-std::unique_ptr<ForLoop> ForLoopNest::AddLoop(absl::string_view suffix,
-                                              llvm::Value* start_index,
-                                              llvm::Value* end_index,
-                                              UnrollMode unroll_mode,
-                                              bool prevent_vectorization) {
+std::unique_ptr<ForLoop> ForLoopNest::AddLoop(
+    absl::string_view suffix, llvm::Value* start_index, llvm::Value* end_index,
+    UnrollMode unroll_mode, bool prevent_vectorization,
+    DExpr expression) {
   return AddLoop(suffix, start_index, end_index, GetConstantWithIndexType(1),
-                 unroll_mode, prevent_vectorization);
+                 unroll_mode, prevent_vectorization, expression);
 }
 
 std::unique_ptr<ForLoop> ForLoopNest::AddLoop(
     absl::string_view suffix, llvm::Value* start_index, llvm::Value* end_index,
-    llvm::Value* stride, UnrollMode unroll_mode, bool prevent_vectorization) {
-  if (inner_loop_body_bb_ != nullptr) {
+    llvm::Value* stride, UnrollMode unroll_mode, bool prevent_vectorization,
+    DExpr expression) {
+  if (inner_loop_body_bb_ != nullptr &&
+      b_->GetInsertBlock() != inner_loop_body_bb_) {
     // Create this loop inside the previous one.
     b_->SetInsertPoint(&*inner_loop_body_bb_->getFirstInsertionPt());
   }
+  llvm::Value* actual_end = end_index;
+  if (expression && expression->is_dynamic()) {
+    // Get batch dim and compare with end_index to use minimum value
+    llvm::Value* expr_value = llvm_ir::EmitExpression(b_, expression);
+    actual_end = b_->CreateSelect(b_->CreateICmpULT(end_index, expr_value),
+                                  end_index, expr_value, "loop_end_min");
+  }
+  // Emitting dynamic expressions may materialize helper instructions into the
+  // current function entry block using a different builder. Re-seat the
+  // insertion point at the end of the current block so the loop preheader is
+  // always in a well-formed state before emitting the loop.
+  llvm::BasicBlock* insert_block = b_->GetInsertBlock();
+  if (insert_block != nullptr && insert_block->getTerminator() == nullptr) {
+    b_->SetInsertPoint(insert_block);
+  }
   std::unique_ptr<ForLoop> loop(new ForLoop(
-      /*prefix=*/name_, suffix, start_index, end_index, stride, unroll_mode,
+      /*prefix=*/name_, suffix, start_index, actual_end, stride, unroll_mode,
       prevent_vectorization));
   loop->Emit(b_);
 
@@ -219,25 +236,46 @@ std::unique_ptr<ForLoop> ForLoopNest::AddLoop(
   return loop;
 }
 
-std::unique_ptr<ForLoop> ForLoopNest::AddLoop(int64_t start_index,
-                                              int64_t end_index,
-                                              absl::string_view suffix,
-                                              UnrollMode unroll_mode,
-                                              bool prevent_vectorization) {
+llvm::Value* ForLoopNest::MaterializeLoopEndValue(int64_t end_index,
+                                                  DExpr expression) {
+  llvm::Value* end = GetConstantWithIndexType(end_index);
+  if (expression && expression->is_dynamic()) {
+    if (inner_loop_body_bb_ != nullptr) {
+      // Keep dynamic loop bound materialization in the same block that will
+      // act as the nested loop preheader so the resulting values dominate the
+      // loop header compare.
+      b_->SetInsertPoint(&*inner_loop_body_bb_->getFirstInsertionPt());
+    }
+    end = EmitExpression(b_, expression);
+  }
+  llvm::BasicBlock* insert_block = b_->GetInsertBlock();
+  if (insert_block != nullptr && insert_block->getTerminator() == nullptr) {
+    b_->SetInsertPoint(insert_block);
+  }
+  return end;
+}
+
+std::unique_ptr<ForLoop> ForLoopNest::AddLoop(
+    int64_t start_index, int64_t end_index, absl::string_view suffix,
+    UnrollMode unroll_mode, bool prevent_vectorization,
+    DExpr expression) {
   CHECK_LE(start_index, end_index);
-  return AddLoop(suffix, GetConstantWithIndexType(start_index),
-                 GetConstantWithIndexType(end_index), unroll_mode,
-                 prevent_vectorization);
+
+  llvm::Value* end = MaterializeLoopEndValue(end_index, expression);
+  return AddLoop(suffix, GetConstantWithIndexType(start_index), end,
+                 unroll_mode, prevent_vectorization);
 }
 
 std::unique_ptr<ForLoop> ForLoopNest::AddLoop(int64_t start_index,
                                               int64_t end_index, int64_t stride,
                                               absl::string_view suffix,
                                               UnrollMode unroll_mode,
-                                              bool prevent_vectorization) {
+                                              bool prevent_vectorization,
+                                              DExpr expression) {
   CHECK_LE(start_index, end_index);
-  return AddLoop(suffix, GetConstantWithIndexType(start_index),
-                 GetConstantWithIndexType(end_index),
+
+  llvm::Value* end = MaterializeLoopEndValue(end_index, expression);
+  return AddLoop(suffix, GetConstantWithIndexType(start_index), end,
                  GetConstantWithIndexType(stride), unroll_mode,
                  prevent_vectorization);
 }
@@ -259,7 +297,9 @@ std::vector<llvm::Value*> ForLoopNest::AddLoopsForShapeOnDimensions(
         /*start_index=*/0,
         /*end_index=*/shape.dimensions(dimension),
         /*suffix=*/
-        llvm_ir::IrName(suffix, absl::StrCat(dimension)));
+        llvm_ir::IrName(suffix, absl::StrCat(dimension)),
+        /*unroll_mode=*/llvm_ir::UnrollMode::kDefaultUnroll,
+        /*prevent_vectorization=*/false, shape.expressions(dimension));
     multi_index[dimension] = loop->GetIndVarValue();
   }
   return multi_index;
diff --git a/third_party/xla/xla/service/llvm_ir/llvm_loop.h b/third_party/xla/xla/service/llvm_ir/llvm_loop.h
index 7aa8ce9e32e950..3a15102bf6dabd 100644
--- a/third_party/xla/xla/service/llvm_ir/llvm_loop.h
+++ b/third_party/xla/xla/service/llvm_ir/llvm_loop.h
@@ -201,14 +201,14 @@ class ForLoopNest {
       absl::string_view suffix, llvm::Value* start_index,
       llvm::Value* end_index, llvm::Value* stride,
       UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll,
-      bool prevent_vectorization = false);
+      bool prevent_vectorization = false, DExpr expression = DExpr());
 
   // Like the above, except that it defaults to a stride of one.
   std::unique_ptr<ForLoop> AddLoop(
       absl::string_view suffix, llvm::Value* start_index,
       llvm::Value* end_index,
       UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll,
-      bool prevent_vectorization = false);
+      bool prevent_vectorization = false, DExpr expression = DExpr());
 
   // A convenient wrapper of the other flavor of AddLoop. The given start and
   // end index are constant.
@@ -216,13 +216,13 @@ class ForLoopNest {
       int64_t start_index, int64_t end_index, int64_t stride,
       absl::string_view suffix,
       UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll,
-      bool prevent_vectorization = false);
+      bool prevent_vectorization = false, DExpr expression = DExpr());
 
   // Like the above, except that it defaults to a stride of one.
   std::unique_ptr<ForLoop> AddLoop(
       int64_t start_index, int64_t end_index, absl::string_view suffix,
       UnrollMode unroll_mode = xla::llvm_ir::UnrollMode::kDefaultUnroll,
-      bool prevent_vectorization = false);
+      bool prevent_vectorization = false, DExpr expression = DExpr());
 
   // Add loops to iterate through the indices within the specified
   // shape. The returned index collects the induction variables of the
@@ -268,6 +268,8 @@ class ForLoopNest {
   llvm::BasicBlock* GetInnerLoopBodyBasicBlock() { return inner_loop_body_bb_; }
 
  private:
+  llvm::Value* MaterializeLoopEndValue(int64_t end_index, DExpr expression);
+
   void SetIndexType(llvm::Type* index_ty) {
     index_type_ = index_ty == nullptr ? b_->getInt64Ty() : index_ty;
   }
diff --git a/third_party/xla/xla/service/llvm_ir/llvm_util.cc b/third_party/xla/xla/service/llvm_ir/llvm_util.cc
index b1db780705b41f..8763a85033624b 100644
--- a/third_party/xla/xla/service/llvm_ir/llvm_util.cc
+++ b/third_party/xla/xla/service/llvm_ir/llvm_util.cc
@@ -68,7 +68,9 @@ limitations under the License.
 #include "xla/layout_util.h"
 #include "xla/literal.h"
 #include "xla/primitive_util.h"
+#include "xla/printer.h"
 #include "xla/service/cpu/cpu_options.h"
+#include "xla/service/cpu/runtime_batch_size.h"
 #include "xla/service/dump.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/llvm_ir/llvm_type_conversion_util.h"
@@ -85,6 +87,8 @@ namespace llvm_ir {
 
 namespace {
 
+constexpr llvm::StringLiteral kBdimValueName("bdim_value");
+
 // This works for most llvm / mlir types. This also accepts a const pointer to
 // objects which have a const print() method.
 template <typename T>
@@ -288,9 +292,11 @@ llvm::Type* ShapeToIrType(const Shape& shape, llvm::LLVMContext& context) {
     result_type =
         llvm::ArrayType::get(result_type, shape.tuple_shapes().size());
   } else if (shape.IsArray()) {
-    for (int64_t dimension : LayoutUtil::MinorToMajor(shape)) {
-      result_type =
-          llvm::ArrayType::get(result_type, shape.dimensions(dimension));
+    auto dimensions = LayoutUtil::MinorToMajor(shape);
+    for (int i = 0; i < dimensions.size(); i++) {
+      bool is_dynamic = shape.expressions(dimensions[i])->is_dynamic();
+      int64_t dim_val = is_dynamic ? 0 : shape.dimensions(dimensions[i]);
+      result_type = llvm::ArrayType::get(result_type, dim_val);
     }
   }
   return result_type;
@@ -811,5 +817,166 @@ void EmitEarlyReturn(llvm::Value* condition, llvm::IRBuilderBase* b,
   b->SetInsertPoint(continued, continued->getFirstInsertionPt());
 }
 
+llvm::Value* GetBatchDimByName(llvm::IRBuilderBase* b, int64_t multiplier,
+                               int64_t offset) {
+  llvm::Function* function = b->GetInsertBlock()->getParent();
+  llvm::LLVMContext& ctx = b->getContext();
+  llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx);
+  llvm::Value* loadedValue = nullptr;
+  llvm::Value* bdim_scaled = nullptr;
+  for (auto& inst : function->getEntryBlock()) {
+    if (inst.getName() == kBdimValueName) {
+      loadedValue = &inst;
+    }
+  }
+  if (!loadedValue) {
+    // If missing, materialize it through the stable CPU runtime interface.
+    llvm::Value* run_options = nullptr;
+    for (auto& arg : function->args()) {
+      if (arg.getName() == "run_options") {
+        run_options = &arg;
+        break;
+      }
+    }
+
+    if (!run_options) {
+      return nullptr;
+    }
+    llvm::IRBuilder<> entry_builder(
+        &function->getEntryBlock(),
+        function->getEntryBlock().getFirstInsertionPt());
+    llvm::FunctionType* getter_type = llvm::FunctionType::get(
+        i64Type, {run_options->getType()}, /*isVarArg=*/false);
+    llvm::FunctionCallee getter = function->getParent()->getOrInsertFunction(
+        xla::cpu::runtime::kGetBatchSizeSymbolName, getter_type);
+    loadedValue = entry_builder.CreateCall(getter, {run_options},
+                                           kBdimValueName);
+  }
+  if (multiplier < 1) {
+    llvm::errs() << "Multiplier is less than 1, this should not happen.\n";
+  } else if (multiplier == 1) {
+    bdim_scaled = loadedValue;
+  } else {
+    llvm::ConstantInt* m = llvm::ConstantInt::get(i64Type, multiplier, true);
+    bdim_scaled = b->CreateMul(loadedValue, m, "bdim_scaled");
+  }
+  if (offset != 0){
+    llvm::ConstantInt* offset_value =
+        llvm::ConstantInt::get(i64Type, offset, true);
+    bdim_scaled = b->CreateAdd(bdim_scaled, offset_value, "bdim_offset");
+  }
+  return bdim_scaled;
+}
+
+static llvm::Value* EmitExpressionImpl(llvm::IRBuilderBase* b,
+                                       const DynExpr& expr) {
+  llvm::LLVMContext& ctx = b->getContext();
+  llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx);
+  if (expr.is_constant()) return llvm::ConstantInt::get(i64Type, expr.get_val(), true);
+  if (expr.kind() == DExpr::Kind::kUnknown) {
+    return nullptr;
+  }
+  if (expr.kind() == DExpr::Kind::kVariable) {
+    // For now we can just use %bdim...
+    return GetBatchDimByName(b);
+  }
+  if (expr.kind() == DExpr::Kind::kMul) {
+    auto* mul_node = static_cast<const Mul*>(&expr);
+    llvm::Value* v_lhs = EmitExpressionImpl(b, *mul_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpressionImpl(b, *mul_node->get_rhs());
+    return b->CreateMul(v_lhs, v_rhs, "mul_dims");
+  }
+  // TODO: Check if this should ever happen
+  if (expr.kind() == DExpr::Kind::kDiv) {
+    auto* div_node = static_cast<const Div*>(&expr);
+    llvm::Value* v_lhs = EmitExpressionImpl(b, *div_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpressionImpl(b, *div_node->get_rhs());
+    return b->CreateSDiv(v_lhs, v_rhs, "div_dims");
+  }
+  if (expr.kind() == DExpr::Kind::kAdd) {
+    auto* add_node = static_cast<const Add*>(&expr);
+    llvm::Value* v_lhs = EmitExpressionImpl(b, *add_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpressionImpl(b, *add_node->get_rhs());
+    return b->CreateAdd(v_lhs, v_rhs, "add_dims");
+  }
+  if (expr.kind() == DExpr::Kind::kSub) {
+    auto* sub_node = static_cast<const Sub*>(&expr);
+    llvm::Value* v_lhs = EmitExpressionImpl(b, *sub_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpressionImpl(b, *sub_node->get_rhs());
+    return b->CreateSub(v_lhs, v_rhs, "sub_dims");
+  }
+  if (expr.kind() == DExpr::Kind::kMax) {
+    auto* max_node = static_cast<const MaxExpr*>(&expr);
+    llvm::Value* v_lhs = EmitExpressionImpl(b, *max_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpressionImpl(b, *max_node->get_rhs());
+    llvm::Value* lhs_is_greater =
+        b->CreateICmpSGT(v_lhs, v_rhs, "max_dims_pred");
+    return b->CreateSelect(lhs_is_greater, v_lhs, v_rhs, "max_dims");
+  }
+  if (expr.kind() == DExpr::Kind::kGt) {
+    auto* gt_node = static_cast<const GtExpr*>(&expr);
+    llvm::Value* v_lhs = EmitExpressionImpl(b, *gt_node->get_lhs());
+    llvm::Value* v_rhs = EmitExpressionImpl(b, *gt_node->get_rhs());
+    llvm::Value* pred = b->CreateICmpSGT(v_lhs, v_rhs, "gt_dims_pred");
+    return b->CreateZExt(pred, i64Type, "gt_dims");
+  }
+  if (expr.kind() == DExpr::Kind::kSelect) {
+    auto* select_node = static_cast<const SelectExpr*>(&expr);
+    llvm::Value* pred = EmitExpressionImpl(b, *select_node->get_pred());
+    llvm::Value* v_true =
+        EmitExpressionImpl(b, *select_node->get_on_true());
+    llvm::Value* v_false =
+        EmitExpressionImpl(b, *select_node->get_on_false());
+    llvm::Value* nonzero = b->CreateICmpNE(
+        pred, llvm::ConstantInt::get(i64Type, 0, true), "select_dims_pred");
+    return b->CreateSelect(nonzero, v_true, v_false, "select_dims");
+  }
+  return nullptr;
+}
+
+llvm::Value* EmitExpression(llvm::IRBuilderBase* b, const DExpr& expr) {
+  if (!expr) return nullptr;
+  DExpr simplified = expr.simplify();
+  StringPrinter printer;
+  simplified->print(&printer);
+  VLOG(2) << "EmitExpression expr=" << std::move(printer).ToString()
+          << " kind=" << static_cast<int>(simplified.kind());
+  llvm::Value* value = EmitExpressionImpl(b, *simplified.get());
+  return value;
+}
+
+llvm::Value* createDynamicGEP(llvm::IRBuilderBase* builder,
+                              llvm::Value* base_ptr,
+                              const std::vector<llvm::Value*>& indices,
+                              absl::Span<const int64_t> dims,
+                              absl::Span<const DExpr> expressions,
+                              llvm::Type* elem_type,
+                              const llvm::Twine& name) {
+  llvm::Value* total_index = builder->getInt64(0);
+  llvm::Type* int64_ty = builder->getInt64Ty();
+
+  for (size_t i = 0; i < indices.size(); ++i) {
+    // The stride is the product of all dimensions to the right of this index.
+    llvm::Value* stride = builder->getInt64(1);
+    for (size_t j = i; j < dims.size(); ++j) {
+      if (expressions[j] && expressions[j]->is_dynamic()) {
+        llvm::Value* expr_value =
+            EmitExpression(builder, expressions[j]);
+        stride = builder->CreateMul(stride, expr_value, "stride.dyn");
+      } else {
+        stride = builder->CreateMul(
+            stride, llvm::ConstantInt::get(int64_ty, dims[j]), "stride.static");
+      }
+    }
+    llvm::Value* scaled_index =
+        builder->CreateMul(indices[i], stride, "idx.scaled");
+    total_index = builder->CreateAdd(total_index, scaled_index, "idx.total");
+  }
+
+  // Final GEP: result = base + total_index * sizeof(elem_type)
+  llvm::Value* gep = builder->CreateGEP(elem_type, base_ptr, total_index, name);
+  return gep;
+}
+
 }  // namespace llvm_ir
 }  // namespace xla
diff --git a/third_party/xla/xla/service/llvm_ir/llvm_util.h b/third_party/xla/xla/service/llvm_ir/llvm_util.h
index 88c1287d2f236d..77bd543c9c0482 100644
--- a/third_party/xla/xla/service/llvm_ir/llvm_util.h
+++ b/third_party/xla/xla/service/llvm_ir/llvm_util.h
@@ -334,6 +334,19 @@ llvm::BasicBlock* EmitReturnBlock(llvm::IRBuilderBase* b);
 void EmitEarlyReturn(llvm::Value* condition, llvm::IRBuilderBase* b,
                      llvm::BasicBlock* return_block = nullptr);
 
+llvm::Value* GetBatchDimByName(llvm::IRBuilderBase* b, int64_t multiplier = 1,
+                               int64_t offset = 0);
+
+llvm::Value* createDynamicGEP(llvm::IRBuilderBase* builder,
+                              llvm::Value* base_ptr,
+                              const std::vector<llvm::Value*>& indices,
+                              absl::Span<const int64_t> dims,
+                              absl::Span<const DExpr> expressions,
+                              llvm::Type* elem_type,
+                              const llvm::Twine& name = "");
+
+llvm::Value* EmitExpression(llvm::IRBuilderBase* b, const DExpr& expr);
+
 }  // namespace llvm_ir
 }  // namespace xla
 
diff --git a/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc b/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc
new file mode 100644
index 00000000000000..4e740b1566e988
--- /dev/null
+++ b/third_party/xla/xla/service/llvm_ir/llvm_util_test.cc
@@ -0,0 +1,83 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/service/llvm_ir/llvm_util.h"
+
+#include <gtest/gtest.h>
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Support/Casting.h"
+#include "xla/service/cpu/runtime_batch_size.h"
+#include "xla/shape_expr.h"
+
+namespace xla {
+namespace llvm_ir {
+namespace {
+
+TEST(LlvmUtilTest, DynamicExpressionDivisionIsSigned) {
+  llvm::LLVMContext context;
+  llvm::Module module("llvm_util_test", context);
+  llvm::IRBuilder<> builder(context);
+  llvm::FunctionType* function_type = llvm::FunctionType::get(
+      llvm::Type::getVoidTy(context), /*isVarArg=*/false);
+  llvm::Function* function = llvm::Function::Create(
+      function_type, llvm::Function::ExternalLinkage, "test", module);
+  llvm::BasicBlock* block =
+      llvm::BasicBlock::Create(context, "entry", function);
+  builder.SetInsertPoint(block);
+
+  llvm::Type* i64_type = builder.getInt64Ty();
+  llvm::Value* batch_dim_address = builder.CreateAlloca(i64_type);
+  builder.CreateLoad(i64_type, batch_dim_address, "bdim_value");
+
+  DExpr expression = (DExpr::Var(1) - 5) / 2;
+  llvm::Value* value = EmitExpression(&builder, expression);
+
+  auto* division = llvm::dyn_cast<llvm::BinaryOperator>(value);
+  ASSERT_NE(division, nullptr);
+  EXPECT_EQ(division->getOpcode(), llvm::Instruction::SDiv);
+}
+
+TEST(LlvmUtilTest, DynamicBatchSizeUsesRuntimeAccessor) {
+  llvm::LLVMContext context;
+  llvm::Module module("llvm_util_test", context);
+  llvm::IRBuilder<> builder(context);
+  llvm::FunctionType* function_type = llvm::FunctionType::get(
+      llvm::Type::getVoidTy(context), {builder.getPtrTy()},
+      /*isVarArg=*/false);
+  llvm::Function* function = llvm::Function::Create(
+      function_type, llvm::Function::ExternalLinkage, "test", module);
+  function->getArg(0)->setName("run_options");
+  llvm::BasicBlock* block =
+      llvm::BasicBlock::Create(context, "entry", function);
+  builder.SetInsertPoint(block);
+
+  llvm::Value* value = GetBatchDimByName(&builder);
+
+  auto* call = llvm::dyn_cast<llvm::CallInst>(value);
+  ASSERT_NE(call, nullptr);
+  ASSERT_NE(call->getCalledFunction(), nullptr);
+  EXPECT_EQ(call->getCalledFunction()->getName(),
+            xla::cpu::runtime::kGetBatchSizeSymbolName);
+}
+
+}  // namespace
+}  // namespace llvm_ir
+}  // namespace xla
diff --git a/third_party/xla/xla/service/llvm_ir/loop_emitter.cc b/third_party/xla/xla/service/llvm_ir/loop_emitter.cc
index 13f6a67764b346..20c2b9098cec17 100644
--- a/third_party/xla/xla/service/llvm_ir/loop_emitter.cc
+++ b/third_party/xla/xla/service/llvm_ir/loop_emitter.cc
@@ -28,11 +28,13 @@ limitations under the License.
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "xla/layout_util.h"
 #include "xla/service/llvm_ir/ir_array.h"
 #include "xla/service/llvm_ir/llvm_loop.h"
+#include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/platform/errors.h"
@@ -183,6 +185,33 @@ std::vector<IrArray::Index> LoopEmitter::EmitIndexAndSetExitBasicBlock(
 
   ForLoopNest loop_nest(loop_name, b_);
 
+  llvm::LLVMContext& ctx = b_->getContext();
+  llvm::IntegerType* i64Type = llvm::IntegerType::getInt64Ty(ctx);
+
+  llvm::PointerType* ptr = llvm::PointerType::getUnqual(ctx);
+  llvm::StructType* callFrameTy = llvm::StructType::create(
+      "XLA_CPU_KernelArg", ptr, ptr, i64Type, ptr, i64Type);
+
+  std::vector<llvm::Value*> dynamic_dims;
+  for (auto dim : shape_.dimensions()) {
+    dynamic_dims.push_back(llvm::ConstantInt::get(i64Type, dim));
+  }
+
+  bool dynamic = false;
+  for (int i = 0; i < shape_.dimensions_size(); i++) {
+    auto expr = shape_.expressions(i);
+    if (expr && expr->is_dynamic()) {
+      dynamic_dims[i] = xla::llvm_ir::EmitExpression(b_, expr);
+      shape_.set_dynamic_dimension(i, true);
+      dynamic = true;
+    }
+  }
+
+  if (dynamic) {
+    // Assign dynamic batch
+    dynamic_dims_ = dynamic_dims;
+  }
+
   IrArray::Index array_index = dynamic_dims_.empty()
                                    ? EmitStaticIndex(&loop_nest, index_type)
                                    : EmitDynamicIndex(&loop_nest, index_type);
diff --git a/third_party/xla/xla/service/llvm_ir/tuple_ops.cc b/third_party/xla/xla/service/llvm_ir/tuple_ops.cc
index b4b3fa30affbcb..14bf83174ca269 100644
--- a/third_party/xla/xla/service/llvm_ir/tuple_ops.cc
+++ b/third_party/xla/xla/service/llvm_ir/tuple_ops.cc
@@ -104,7 +104,7 @@ llvm::Value* EmitGetTupleElement(const Shape& target_shape, int64_t index,
   llvm::LoadInst* src_buffer = b->CreateLoad(element_pointee_type, element_ptr);
 
   // Mark the loaded pointer as dereferenceable if we know its shape.
-  if (!target_shape.IsOpaque()) {
+  if (!target_shape.IsOpaque() && !target_shape.has_dynamic_expr()) {
     SetDereferenceableMetadataForLoad(
         src_buffer,
         ByteSizeOf(target_shape, src_buffer->getModule()->getDataLayout()));
diff --git a/third_party/xla/xla/service/reduce_scatter_combiner.cc b/third_party/xla/xla/service/reduce_scatter_combiner.cc
index a86271036f6baa..361629b34a6d5e 100644
--- a/third_party/xla/xla/service/reduce_scatter_combiner.cc
+++ b/third_party/xla/xla/service/reduce_scatter_combiner.cc
@@ -131,9 +131,9 @@ absl::Status CombineReduceScatters(
       std::swap((*perm)[most_frequent_dim], (*perm)[rs->scatter_dimension()]);
 
       // Bitcast operand and update output shape.
-      operands.back() =
-          computation.AddInstruction(HloInstruction::CreateBitcast(
-              ShapeUtil::PermuteDimensions(*perm, operand_shape), operand));
+      auto sh = ShapeUtil::PermuteDimensions(*perm, operand_shape);
+      operands.back() = computation.AddInstruction(
+          HloInstruction::CreateBitcast(sh, operand));
       output_shapes.back() = ShapeUtil::PermuteDimensions(*perm, hlo->shape());
     }
   }
diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc
index 7985c930d812b5..ce006f2a132dd3 100644
--- a/third_party/xla/xla/service/shape_inference.cc
+++ b/third_party/xla/xla/service/shape_inference.cc
@@ -77,6 +77,15 @@ bool CompatibleDimensionSizes(int64_t size_a, int64_t size_b) {
          size_a == size_b;
 }
 
+DExpr SymbolicElementsIn(const Shape& shape) {
+  DExpr product = DExpr::Const(1);
+  for (int64_t i = 0; i < shape.dimensions_size(); ++i) {
+    const DExpr& expr = shape.expressions(i);
+    product = product * (expr ? expr : DExpr::Const(shape.dimensions(i)));
+  }
+  return product.simplify();
+}
+
 absl::Status ExpectArray(const Shape& shape, absl::string_view op_type) {
   if (!shape.IsArray()) {
     return InvalidArgument("Expected array argument for %s, but got %s.",
@@ -195,6 +204,7 @@ absl::StatusOr<Shape> InferWindowOutputShape(const Shape& base_shape,
 
   std::vector<int64_t> output_dimensions(window.dimensions_size());
   std::vector<bool> output_is_dynamic(window.dimensions_size());
+  std::vector<DExpr> output_expressions(window.dimensions_size());
   for (int64_t i = 0; i < window.dimensions_size(); ++i) {
     const auto& dim = window.dimensions(i);
     if (dim.size() <= 0) {
@@ -216,24 +226,45 @@ absl::StatusOr<Shape> InferWindowOutputShape(const Shape& base_shape,
           window.DebugString());
     }
 
-    if (IsUnboundedDynamicSize(ShapeUtil::GetDimension(base_shape, i))) {
+    const int64_t input_dimension = ShapeUtil::GetDimension(base_shape, i);
+    const DExpr& input_expression = base_shape.expressions(i);
+    const int64_t dilated_window =
+        window_util::DilatedBound(dim.size(), dim.window_dilation());
+    if (IsUnboundedDynamicSize(input_dimension)) {
       output_dimensions[i] = Shape::kUnboundedSize;
     } else {
       const int64_t dilated_base = window_util::DilatedBound(
-          ShapeUtil::GetDimension(base_shape, i), dim.base_dilation());
+          input_dimension, dim.base_dilation());
       const int64_t padded_dilated_base =
           dim.padding_low() + dilated_base + dim.padding_high();
-      const int64_t dilated_window =
-          window_util::DilatedBound(dim.size(), dim.window_dilation());
-
       output_dimensions[i] = window_util::StridedBound(
           padded_dilated_base, dilated_window, dim.stride());
     }
     output_is_dynamic[i] = base_shape.is_dynamic_dimension(i);
+    if (input_expression && input_expression->is_constant()) {
+      output_expressions[i] = DExpr::Const(output_dimensions[i]);
+      continue;
+    }
+
+    DExpr dilated_base_expr = input_expression;
+    if (dim.base_dilation() != 1) {
+      dilated_base_expr =
+          DExpr::Max((dim.base_dilation() * (input_expression - 1)) + 1,
+                     DExpr::Const(0))
+              .simplify();
+    }
+    DExpr padded_dilated_base_expr =
+        (dilated_base_expr + dim.padding_low() + dim.padding_high()).simplify();
+    DExpr strided_bound_expr =
+        (padded_dilated_base_expr - dilated_window + 1 + dim.stride() - 1) /
+        dim.stride();
+    output_expressions[i] =
+        DExpr::Max(strided_bound_expr, DExpr::Const(0)).simplify();
   }
 
   return ShapeUtil::MakeValidatedShape(element_type, output_dimensions,
-                                       output_is_dynamic);
+                                       output_is_dynamic,
+                                       output_expressions);
 }
 
 // Encapsulates inferred dimension size and bound size.
@@ -472,6 +503,7 @@ absl::StatusOr<DimAndBound> InferMostSpecificDimAndBound(int64_t dim,
   int64_t last_dim = operand_shape.dimensions_size() - 1;
   std::vector<bool> is_dynamic(operand_shape.dimensions_size());
   std::vector<int64_t> dimensions(operand_shape.dimensions_size());
+  std::vector<xla::DExpr> expressions(operand_shape.dimensions_size());
 
   TF_RET_CHECK(operand_shape.dimensions(last_dim) >= k)
       << "k=" << k << " is larger than the last dimension of size="
@@ -480,10 +512,13 @@ absl::StatusOr<DimAndBound> InferMostSpecificDimAndBound(int64_t dim,
     is_dynamic[i] =
         i == last_dim ? false : operand_shape.is_dynamic_dimension(i);
     dimensions[i] = i == last_dim ? k : operand_shape.dimensions(i);
+    expressions[i] =
+        i == last_dim ? xla::DExpr::Const(k) : operand_shape.expressions(i);
   }
 
-  Shape out = ShapeUtil::MakeShape(operand_shape.element_type(), dimensions,
-                                   is_dynamic);
+  Shape out =
+      ShapeUtil::MakeShape(operand_shape.element_type(), dimensions, is_dynamic,
+                           expressions);
   Shape idxs_shape = ShapeUtil::ChangeElementType(out, PrimitiveType::S32);
   return ShapeUtil::MakeTupleShape({out, idxs_shape});
 }
@@ -542,9 +577,11 @@ absl::StatusOr<DimAndBound> InferMostSpecificDimAndBound(int64_t dim,
   int64_t rank = arg_shape->dimensions_size();
   std::vector<int64_t> inferred_sizes(rank, Shape::kUnboundedSize);
   std::vector<int64_t> inferred_bounds(rank, Shape::kUnboundedSize);
+  std::vector<xla::DExpr> inferred_expressions(rank, xla::DExpr::Const(0));
   // Note: for the concatenate dimension, 0 should be the identity element:
   // Any dim size can keep unchanged when concatenated with 0
   inferred_sizes[dimension] = 0;
+  inferred_expressions[dimension] = xla::DExpr::Const(0);
 
   for (const Shape* shape : arg_shapes) {
     for (int dim = 0; dim < rank; ++dim) {
@@ -554,24 +591,32 @@ absl::StatusOr<DimAndBound> InferMostSpecificDimAndBound(int64_t dim,
       int64_t leftSize = inferred_sizes[dim];
       int64_t rightSize = dimension_size;
       int64_t leftBound = inferred_bounds[dim];
+      const xla::DExpr& left_expression = inferred_expressions[dim];
       int64_t rightBound = shape->is_dynamic_dimension(dim)
                                ? dimension_size
                                : Shape::kUnboundedSize;
+      const xla::DExpr& right_expression = shape->expressions(dim);
+      xla::DExpr inferred_expression = xla::DExpr::Const(0);
+
       if (dim == dimension) {
         inferred_dim_and_bound = InferConcatenatedDimAndBound(
             leftSize, rightSize, leftBound, rightBound);
+        inferred_expression = left_expression + right_expression;
       } else {
         TF_ASSIGN_OR_RETURN(
             inferred_dim_and_bound,
             InferMostSpecificDimAndBound(dim, leftSize, rightSize, leftBound,
                                          rightBound));
+        inferred_expression = right_expression;
       }
       inferred_sizes[dim] = inferred_dim_and_bound.dimension;
       inferred_bounds[dim] = inferred_dim_and_bound.bound;
+      inferred_expressions[dim] = inferred_expression;
     }
   }
 
-  Shape result = ShapeUtil::MakeShape(element_type, inferred_sizes);
+  Shape result =
+      ShapeUtil::MakeShape(element_type, inferred_sizes, inferred_expressions);
   for (int64_t i = 0; i < inferred_bounds.size(); ++i) {
     if (!IsUnboundedDynamicSize(inferred_bounds[i]) ||
         IsUnboundedDynamicSize(inferred_sizes[i])) {
@@ -756,6 +801,7 @@ absl::StatusOr<DimAndBound> InferMostSpecificDimAndBound(int64_t dim,
 
   std::vector<int64_t> dimensions(operand_shape.dimensions_size());
   std::vector<bool> is_dynamic(operand_shape.dimensions_size());
+  std::vector<DExpr> expressions(operand_shape.dimensions_size());
   for (int64_t i = 0; i < operand_shape.dimensions_size(); ++i) {
     const auto& p = padding_config.dimensions(i);
     if (operand_shape.is_unbounded_dynamic_dimension(i)) {
@@ -771,11 +817,13 @@ absl::StatusOr<DimAndBound> InferMostSpecificDimAndBound(int64_t dim,
       }
     }
     is_dynamic[i] = operand_shape.is_dynamic_dimension(i);
+    auto diff = dimensions[i] - operand_shape.dimensions(i);
+    expressions[i] = operand_shape.expressions(i) + diff;
   }
 
   return ShapeUtil::MakeShape(
       ShapeUtil::HigherPrecisionElementType(operand_shape, padding_value_shape),
-      dimensions, is_dynamic);
+      dimensions, is_dynamic, expressions);
 }
 
 // Current DotDimensionNumbers Requirements:
@@ -920,7 +968,9 @@ absl::Status CheckDotDimensionConstraints(
 void GenerateDotResultDimensions(
     const Shape& lhs, const Shape& rhs,
     const DotDimensionNumbers& dimension_numbers,
-    std::vector<int64_t>& dimensions, std::vector<bool>& is_dynamic,
+    std::vector<int64_t>& dimensions,
+    std::vector<DExpr>& expressions,
+    std::vector<bool>& is_dynamic,
     std::vector<int64_t> rhs_group_dimensions = {}) {
   const auto& lhs_batch_dimensions = dimension_numbers.lhs_batch_dimensions();
   const auto lhs_batch_dimensions_size =
@@ -930,9 +980,11 @@ void GenerateDotResultDimensions(
       dimension_numbers.rhs_contracting_dimensions().size() -
       dimension_numbers.rhs_batch_dimensions().size();
   dimensions.reserve(lhs_batch_dimensions_size);
+  expressions.reserve(lhs_batch_dimensions_size);
   is_dynamic.reserve(lhs_batch_dimensions_size);
   for (const int64_t lhs_dim : lhs_batch_dimensions) {
     dimensions.push_back(lhs.dimensions(lhs_dim));
+    expressions.push_back(lhs.expressions(lhs_dim));
     is_dynamic.push_back(lhs.is_dynamic_dimension(lhs_dim));
   }
   for (int64_t i = 0; i < lhs.dimensions_size(); i++) {
@@ -940,6 +992,7 @@ void GenerateDotResultDimensions(
                                i) &&
         !absl::c_linear_search(dimension_numbers.lhs_batch_dimensions(), i)) {
       dimensions.push_back(lhs.dimensions(i));
+      expressions.push_back(lhs.expressions(i));
       is_dynamic.push_back(lhs.is_dynamic_dimension(i));
     }
   }
@@ -949,6 +1002,7 @@ void GenerateDotResultDimensions(
         !absl::c_linear_search(dimension_numbers.rhs_batch_dimensions(), i) &&
         !absl::c_linear_search(rhs_group_dimensions, i)) {
       dimensions.push_back(rhs.dimensions(i));
+      expressions.push_back(rhs.expressions(i));
       is_dynamic.push_back(rhs.is_dynamic_dimension(i));
     }
   }
@@ -990,12 +1044,14 @@ void GenerateDotResultDimensions(
 
   std::vector<int64_t> dimensions;
   std::vector<bool> is_dynamic;
+  std::vector<DExpr> expressions;
   GenerateDotResultDimensions(lhs, rhs, dimension_numbers, dimensions,
-                              is_dynamic);
+                              expressions, is_dynamic);
 
   PrimitiveType type = preferred_element_type.value_or(
       ShapeUtil::HigherPrecisionElementType(lhs, rhs));
-  Shape result = ShapeUtil::MakeShape(type, dimensions, is_dynamic);
+  Shape result =
+      ShapeUtil::MakeShape(type, dimensions, is_dynamic, expressions);
 
   TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(result));
   VLOG(2) << "inferred dot shape: " << ShapeUtil::HumanString(result);
@@ -1193,16 +1249,24 @@ void GenerateDotResultDimensions(
   PrimitiveType type = preferred_element_type.value_or(
       ShapeUtil::HigherPrecisionElementType(lhs, rhs));
   std::vector<int64_t> dimensions;
+  std::vector<DExpr> expressions;
   std::vector<bool> is_dynamic;
   // Add the group dimension to the result shape in case of ragged contracting.
   if (mode == kContracting) {
     dimensions.push_back(num_groups);
     is_dynamic.push_back(is_dynamic_group_sizes);
+    const DExpr& group_expr =
+        group_sizes.expressions(group_sizes.dimensions_size() - 1);
+    expressions.push_back(group_expr ? group_expr
+                                     : (is_dynamic_group_sizes
+                                            ? DExpr::Unknown(70)
+                                            : DExpr::Const(num_groups)));
   }
   GenerateDotResultDimensions(lhs, rhs, dimension_numbers, dimensions,
-                              is_dynamic, rhs_group_dimensions);
+                              expressions, is_dynamic, rhs_group_dimensions);
 
-  Shape result = ShapeUtil::MakeShape(type, dimensions, is_dynamic);
+  Shape result =
+      ShapeUtil::MakeShape(type, dimensions, is_dynamic, expressions);
   TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(result));
   VLOG(2) << "inferred ragged dot shape: " << ShapeUtil::HumanString(result);
   return result;
@@ -1248,12 +1312,15 @@ void GenerateDotResultDimensions(
   // Build the resulting shape dimensions.
   std::vector<int64_t> dimensions;
   std::vector<bool> is_dynamic;
+  std::vector<DExpr> expressions;
   for (int64_t i = 0; i < operand_shape.dimensions_size(); ++i) {
     dimensions.push_back(i != sparsity.dimension() ? operand_shape.dimensions(i)
                                                    : metadata_dimension_size);
     is_dynamic.push_back(operand_shape.is_dynamic_dimension(i));
+    expressions.push_back(operand_shape.expressions(i));
   }
-  return ShapeUtil::MakeShape(element_type, dimensions, is_dynamic);
+  return ShapeUtil::MakeShape(element_type, dimensions, is_dynamic,
+                              expressions);
 }
 
 /* static */ absl::StatusOr<Shape>
@@ -1267,6 +1334,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs,
   // from the lhs/rhs pair in every index.
   std::vector<int64_t> output_dimensions(lhs.dimensions_size());
   std::vector<bool> output_dimensions_is_dynamic(lhs.dimensions_size());
+  std::vector<DExpr> output_dimensions_expressions(lhs.dimensions_size());
   for (int64_t i = 0; i < lhs.dimensions_size(); ++i) {
     if (lhs.dimensions(i) == 1 || rhs.dimensions(i) == 1) {
       // For the unbounded case, the operand with 1 should be broadcasted to the
@@ -1283,6 +1351,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs,
       output_dimensions_is_dynamic[i] = lhs.dimensions(i) == 1
                                             ? rhs.is_dynamic_dimension(i)
                                             : lhs.is_dynamic_dimension(i);
+      output_dimensions_expressions[i] = lhs.dimensions(i) == 1
+                                             ? rhs.expressions(i)
+                                             : lhs.expressions(i);
     } else if (lhs.dimensions(i) == rhs.dimensions(i)) {
       // LHS | RHS | Result
       // X   | X   | X
@@ -1293,6 +1364,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs,
       output_dimensions[i] = lhs.dimensions(i);
       output_dimensions_is_dynamic[i] =
           lhs.is_dynamic_dimension(i) || rhs.is_dynamic_dimension(i);
+      output_dimensions_expressions[i] = lhs.expressions(i);
     } else if (lhs.is_unbounded_dynamic_dimension(i) ||
                rhs.is_unbounded_dynamic_dimension(i)) {
       // For the last two rows, consider when <=X turns out to be 1 and ? turns
@@ -1309,6 +1381,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs,
       output_dimensions_is_dynamic[i] = lhs.is_unbounded_dynamic_dimension(i)
                                             ? rhs.is_dynamic_dimension(i)
                                             : lhs.is_dynamic_dimension(i);
+      output_dimensions_expressions[i] = lhs.is_unbounded_dynamic_dimension(i)
+                                             ? rhs.expressions(i)
+                                             : lhs.expressions(i);
     } else {
       return InvalidArgument("Binary op with incompatible shapes: %s and %s.",
                              ShapeUtil::HumanString(lhs),
@@ -1317,7 +1392,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs,
   }
 
   return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs),
-                              output_dimensions, output_dimensions_is_dynamic);
+                              output_dimensions, output_dimensions_is_dynamic,
+                              output_dimensions_expressions);
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferInDimBroadcastShape(
@@ -1398,6 +1474,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs,
                     dimension_to_match, larger_shape.dimensions_size()));
     }
     int64_t small_dimension_size = smaller_shape.dimensions(i);
+    const DExpr& small_dimension_exp = smaller_shape.expressions(i);
     int64_t large_dimension_size = larger_shape.dimensions(dimension_to_match);
     bool small_is_dynamic = smaller_shape.is_dynamic_dimension(i);
     bool large_is_dynamic =
@@ -1436,6 +1513,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(const Shape& lhs,
 
     output_shape.set_dimensions(dimension_to_match, small_dimension_size,
                                 small_is_dynamic);
+    output_shape.set_expression(dimension_to_match, small_dimension_exp);
   }
 
   return output_shape;
@@ -1484,6 +1562,7 @@ ShapeInference::InferElementwiseBinaryOpShape(
     for (int64_t i = 0; i < rhs.dimensions_size(); ++i) {
       if (rhs.is_dynamic_dimension(i)) {
         result.set_dynamic_dimension(i, true);
+        result.set_expression(i, rhs.expressions(i));
       }
     }
 
@@ -1779,7 +1858,9 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
       output_shape.element_type(), arg_shape->dimensions(),
       /*dynamic_dimensions=*/
       std::vector<bool>(arg_shape->dynamic_dimensions().begin(),
-                        arg_shape->dynamic_dimensions().end()));
+                        arg_shape->dynamic_dimensions().end()),
+      /*expressions=*/
+      arg_shape->expressions());
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferBatchNormTrainingShape(
@@ -1864,8 +1945,12 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
 
   const int64_t feature_count = operand_shape.dimensions(feature_index);
   bool dynamic_feature = operand_shape.is_dynamic_dimension(feature_index);
-  Shape output_shape_for_mean_and_var = ShapeUtil::MakeShape(
-      operand_shape.element_type(), {feature_count}, {dynamic_feature});
+  const DExpr& expression_feature = operand_shape.expressions(feature_index);
+  std::array<DExpr, 1> feature_expressions = {expression_feature};
+
+  Shape output_shape_for_mean_and_var =
+      ShapeUtil::MakeShape(operand_shape.element_type(), {feature_count},
+                           {dynamic_feature}, feature_expressions);
 
   if (!CompatibleDimensionSizes(ShapeUtil::GetDimension(offset_shape, 0),
                                 feature_count)) {
@@ -2148,8 +2233,12 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
 
   const int64_t feature_count = operand_shape.dimensions(feature_index);
   bool dynamic_feature = operand_shape.is_dynamic_dimension(feature_index);
+  const DExpr& expression_feature = operand_shape.expressions(feature_index);
+  std::array<DExpr, 1> feature_expressions = {expression_feature};
+
   Shape feature_shape = ShapeUtil::MakeShape(
-      operand_shape.element_type(), {feature_count}, {dynamic_feature});
+      operand_shape.element_type(), {feature_count}, {dynamic_feature},
+      feature_expressions);
 
   if (!CompatibleDimensionSizes(ShapeUtil::GetDimension(mean_shape, 0),
                                 feature_count)) {
@@ -2398,13 +2487,18 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
   }
 
   std::vector<bool> dynamic_dimensions(input_spatial_dims.size());
-  for (auto it = input_spatial_dims.begin(); it != input_spatial_dims.end();
-       ++it) {
-    dynamic_dimensions[it - input_spatial_dims.begin()] =
-        IsUnboundedDynamicSize(*it);
+  std::vector<DExpr> expressions(input_spatial_dims.size());
+  for (int i = 0; i < input_spatial_dims.size(); ++i) {
+    const int64_t input_spatial_dimension =
+        dnums.input_spatial_dimensions(i);
+    dynamic_dimensions[i] = IsUnboundedDynamicSize(input_spatial_dims[i]);
+    expressions[i] = lhs.expressions(input_spatial_dimension)
+                         ? lhs.expressions(input_spatial_dimension)
+                         : DExpr::Const(input_spatial_dims[i]);
   }
   Shape base_shape = ShapeUtil::MakeShape(
-      lhs.element_type(), input_spatial_dims, dynamic_dimensions);
+      lhs.element_type(), input_spatial_dims, dynamic_dimensions,
+      expressions);
   TF_ASSIGN_OR_RETURN(
       Shape window_output_shape,
       InferWindowOutputShape(base_shape, window, lhs.element_type()));
@@ -2418,6 +2512,15 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
         window_output_shape.dimensions(i);
   }
   std::vector<bool> is_dynamic(num_dims);
+  std::vector<DExpr> output_expressions(num_dims);
+  output_expressions[dnums.output_batch_dimension()] =
+      lhs.expressions(dnums.input_batch_dimension()) / batch_group_count;
+  output_expressions[dnums.output_feature_dimension()] =
+      DExpr::Const(kernel_output_features);
+  for (int i = 0; i < num_spatial_dims; ++i) {
+    output_expressions[dnums.output_spatial_dimensions(i)] =
+        window_output_shape.expressions(i);
+  }
   for (int i = 0; i < num_dims; i++) {
     if (lhs.is_dynamic_dimension(i)) {
       if (i == dnums.input_batch_dimension()) {
@@ -2457,7 +2560,8 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
   }
   PrimitiveType type = preferred_element_type.value_or(
       ShapeUtil::HigherPrecisionElementType(lhs, rhs));
-  return ShapeUtil::MakeShape(type, dimensions, is_dynamic);
+  return ShapeUtil::MakeShape(type, dimensions, is_dynamic,
+                              output_expressions);
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferFftShape(
@@ -2780,8 +2884,10 @@ ShapeInference::InferScalarBroadcastShape(absl::Span<const Shape> shapes) {
 
   const std::vector<bool> dynamic_dimensions(shape.dynamic_dimensions().begin(),
                                              shape.dynamic_dimensions().end());
+  auto exprs = shape.expressions();
+  std::vector<DExpr> expressions(exprs.begin(), exprs.end());
   return ShapeUtil::MakeShape(shape.element_type(), new_dimensions,
-                              dynamic_dimensions);
+                              dynamic_dimensions, expressions);
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferAllToAllTupleShape(
@@ -2923,23 +3029,28 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
 
   std::vector<int64_t> new_dimensions;
   std::vector<bool> new_is_dynamic;
+  std::vector<DExpr> new_expressions;
   for (int i = 0; i < arg.dimensions_size(); ++i) {
     if (dimensions_to_reduce_set.find(i) == dimensions_to_reduce_set.end()) {
       new_dimensions.push_back(arg.dimensions(i));
       new_is_dynamic.push_back(arg.is_dynamic_dimension(i));
+      new_expressions.push_back(arg.expressions(i));
     }
   }
 
   if (ShapeUtil::IsScalar(to_apply.result())) {
-    return ShapeUtil::MakeShape(to_apply.result().element_type(),
-                                new_dimensions, new_is_dynamic);
+    return ShapeUtil::MakeShape(
+        to_apply.result().element_type(), new_dimensions, new_is_dynamic,
+        new_expressions);
   } else {
     std::vector<Shape> result_subshapes;
     const auto& tuple_shapes = to_apply.result().tuple_shapes();
     result_subshapes.reserve(tuple_shapes.size());
     for (const Shape& subshape : tuple_shapes) {
-      result_subshapes.push_back(ShapeUtil::MakeShape(
-          subshape.element_type(), new_dimensions, new_is_dynamic));
+      auto new_shape = ShapeUtil::MakeShape(
+          subshape.element_type(), new_dimensions, new_is_dynamic,
+          new_expressions);
+      result_subshapes.push_back(new_shape);
     }
     return ShapeUtil::MakeTupleShape(result_subshapes);
   }
@@ -3172,7 +3283,9 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferSliceShape(
     const Shape& arg, absl::Span<const int64_t> starts,
-    absl::Span<const int64_t> limits, absl::Span<const int64_t> strides) {
+    absl::Span<const int64_t> limits, absl::Span<const int64_t> strides,
+    absl::Span<const DExpr> start_exprs,
+    absl::Span<const DExpr> limit_exprs) {
   auto error = [&](const std::string& message) {
     return InvalidArgument(
         "%s in slice operation; argument shape: %s; starts: {%s}; limits: "
@@ -3202,8 +3315,10 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
   }
 
   std::vector<int64_t> sizes;
+  std::vector<DExpr> expressions;
   const auto starts_size = starts.size();
   sizes.reserve(starts_size);
+  expressions.reserve(starts_size);
   for (int64_t dimension = 0; dimension < starts_size; ++dimension) {
     int64_t start_index = starts[dimension];
     int64_t limit_index = limits[dimension];
@@ -3231,6 +3346,15 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
       return InvalidArgument("Stride (%d) must be positive.", stride);
     }
     sizes.push_back((limit_index - start_index + stride - 1) / stride);
+
+    DExpr limit_expr =
+        limit_exprs.empty() ? DExpr::Const(limit_index) : limit_exprs[dimension];
+    DExpr start_expr =
+        start_exprs.empty() ? DExpr::Const(start_index) : start_exprs[dimension];
+
+    auto new_expr =
+        limit_expr - start_expr + DExpr::Const(stride) - DExpr::Const(1);
+    expressions.push_back((new_expr / DExpr::Const(stride)).simplify());
   }
 
   std::vector<bool> is_dynamic(arg.dimensions_size());
@@ -3242,12 +3366,14 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
     is_dynamic[i] = arg.is_bounded_dynamic_dimension(i);
   }
 
-  return ShapeUtil::MakeShape(arg.element_type(), sizes, is_dynamic);
+  return ShapeUtil::MakeShape(arg.element_type(), sizes, is_dynamic,
+                              expressions);
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferDynamicSliceShape(
     const Shape& operand_shape, absl::Span<const Shape> start_index_shapes,
-    absl::Span<const int64_t> slice_sizes, bool allow_scalar_indices) {
+    absl::Span<const int64_t> slice_sizes,
+    absl::Span<const DExpr> slice_exprs, bool allow_scalar_indices) {
   TF_RETURN_IF_ERROR(ExpectArray(operand_shape, "operand of dynamic slice"));
   auto number_of_indices = start_index_shapes.size();
   // TODO(b/118437727): Remove this path.
@@ -3346,8 +3472,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
     VLOG(2) << StrFormat("slice_sizes[%d] = %d", dim, slice_dim_size);
   }
 
-  Shape result =
-      ShapeUtil::MakeShape(operand_shape.element_type(), slice_sizes);
+  Shape result = ShapeUtil::MakeShape(operand_shape.element_type(), slice_sizes,
+                                      slice_exprs);
 
   for (int64_t dimension = 0; dimension < operand_shape.dimensions_size();
        ++dimension) {
@@ -3646,7 +3772,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
 }
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferBroadcastShape(
-    const Shape& operand, absl::Span<const int64_t> broadcast_sizes) {
+    const Shape& operand, absl::Span<const int64_t> broadcast_sizes,
+    absl::Span<const DExpr> broadcast_exprs) {
   // This method is used to infer shape for xla::BroadcastInDim.
   TF_RETURN_IF_ERROR(ExpectArray(operand, "operand of broadcast"));
   TF_RET_CHECK(!operand.is_unbounded_dynamic());
@@ -3666,12 +3793,14 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
   std::copy(operand.dimensions().begin(), operand.dimensions().end(),
             dimensions.begin() + broadcast_sizes.size());
 
-  TF_ASSIGN_OR_RETURN(Shape result, ShapeUtil::MakeValidatedShape(
-                                        operand.element_type(), dimensions));
+  TF_ASSIGN_OR_RETURN(
+      Shape result, ShapeUtil::MakeValidatedShape(operand.element_type(),
+                                                  dimensions, broadcast_exprs));
   for (int64_t i = 0; i < operand.dimensions_size(); ++i) {
     result.set_dynamic_dimension(broadcast_sizes.size() + i,
                                  operand.is_dynamic_dimension(i));
   }
+
   return result;
 }
 
@@ -3735,7 +3864,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
 /* static */ absl::StatusOr<Shape> ShapeInference::InferDynamicReshapeShape(
     const Shape& operand, absl::Span<const Shape* const> dim_size_shapes,
     absl::Span<const int64_t> new_size_bounds,
-    const std::vector<bool>& dims_are_dynamic) {
+    const std::vector<bool>& dims_are_dynamic,
+    absl::Span<const DExpr> expressions) {
   if (new_size_bounds.size() != dims_are_dynamic.size()) {
     return InvalidArgument(
         "DynamicReshape has to have the same number of elements in new_sizes "
@@ -3751,9 +3881,8 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
           dim_size_shape->ToString());
     }
   }
-
   Shape inferred_shape = ShapeUtil::MakeShape(
-      operand.element_type(), new_size_bounds, dims_are_dynamic);
+      operand.element_type(), new_size_bounds, dims_are_dynamic, expressions);
   if (ShapeUtil::ElementsIn(operand) != ShapeUtil::ElementsIn(inferred_shape)) {
     return InvalidArgument(
         "Reshape operation has mismatched element counts: from=%d (%s) "
@@ -3767,10 +3896,21 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
 
 /* static */ absl::StatusOr<Shape> ShapeInference::InferReshapeShape(
     const Shape& operand, absl::Span<const int64_t> dimensions,
-    int64_t inferred_dimension) {
+    int64_t inferred_dimension, absl::Span<const DExpr> expressions) {
   TF_RETURN_IF_ERROR(ExpectArray(operand, "reshape"));
   Shape inferred_shape =
-      ShapeUtil::MakeShape(operand.element_type(), dimensions);
+      ShapeUtil::MakeShape(operand.element_type(), dimensions, expressions);
+
+  if (expressions.empty() && operand.expressions().size() > 0 &&
+      operand.expressions(0)->is_dynamic()) {
+    return InvalidArgument("Expressions is empty but operand is dynamic");
+  }
+
+  // if (!expressions.empty() && expressions[0]->is_constant() &&
+  //     expressions[0]->get_val() == 977) {
+  //   return InvalidArgument("Expressions[0] is the magic number (977).");
+  // }
+
   VLOG(3) << "Reshape inferred shape: "
           << ShapeUtil::HumanString(inferred_shape);
 
@@ -3790,6 +3930,16 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
         ShapeUtil::ElementsIn(inferred_shape),
         ShapeUtil::HumanString(inferred_shape));
   }
+  if (!expressions.empty()) {
+    DExpr input_elements = SymbolicElementsIn(operand);
+    DExpr output_elements = SymbolicElementsIn(inferred_shape);
+    if (!DynExpr::equal(input_elements.get(), output_elements.get())) {
+      return InvalidArgument(
+          "Reshape operation has mismatched symbolic element counts: "
+          "from=%s to=%s.",
+          ShapeUtil::HumanString(operand), ShapeUtil::HumanString(inferred_shape));
+    }
+  }
 
   std::vector<int64_t> indices(operand.dimensions_size());
   std::iota(indices.begin(), indices.end(), 0);
@@ -4017,6 +4167,22 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) {
                          on_true.is_dynamic_dimension(dimension) ||
                          on_false.is_dynamic_dimension(dimension));
     }
+    const DExpr& on_true_expr = on_true.expressions(dimension);
+    const DExpr& on_false_expr = on_false.expressions(dimension);
+    const bool has_dynamic_expression =
+        (on_true_expr && on_true_expr->is_dynamic()) ||
+        (on_false_expr && on_false_expr->is_dynamic());
+    if (has_dynamic_expression) {
+      if (!on_true_expr || !on_false_expr ||
+          !DynExpr::equal(on_true_expr, on_false_expr)) {
+        return InvalidArgument(
+            "Select operands have mismatched expressions in dimension %d: "
+            "on_true=%s, on_false=%s.",
+            dimension, ShapeUtil::HumanString(on_true),
+            ShapeUtil::HumanString(on_false));
+      }
+      result.set_expression(dimension, on_true_expr);
+    }
   }
   if (result.has_layout()) {
     result.mutable_layout()->set_element_size_in_bits(
@@ -4250,18 +4416,25 @@ static absl::Status ValidateGatherDimensionNumbers(
   std::vector<int64_t> expanded_start_indices_shape;
   // Also tracks if an output dimension is dynamic.
   std::vector<bool> expanded_start_indices_shape_dynamic_dimensions;
+  std::vector<DExpr> expanded_start_indices_shape_expressions;
   expanded_start_indices_shape.reserve(start_indices_shape.dimensions_size());
   expanded_start_indices_shape_dynamic_dimensions.reserve(
       start_indices_shape.dimensions_size());
+  expanded_start_indices_shape_expressions.reserve(
+      start_indices_shape.dimensions_size());
   absl::c_copy(start_indices_shape.dimensions(),
                std::back_inserter(expanded_start_indices_shape));
   absl::c_copy(
       start_indices_shape.dynamic_dimensions(),
       std::back_inserter(expanded_start_indices_shape_dynamic_dimensions));
+  absl::c_copy(
+      start_indices_shape.expressions(),
+      std::back_inserter(expanded_start_indices_shape_expressions));
   if (expanded_start_indices_shape.size() ==
       gather_dim_numbers.index_vector_dim()) {
     expanded_start_indices_shape.push_back(1);
     expanded_start_indices_shape_dynamic_dimensions.push_back(false);
+    expanded_start_indices_shape_expressions.push_back(DExpr::Const(1));
   }
 
   TF_RETURN_IF_ERROR(ValidateGatherDimensionNumbers(
@@ -4328,10 +4501,12 @@ static absl::Status ValidateGatherDimensionNumbers(
   output_dim_bounds.reserve(result_rank);
 
   std::vector<bool> output_dim_is_dynamic;
+  std::vector<DExpr> output_expressions;
   output_dim_is_dynamic.reserve(result_rank);
   for (int64_t i = 0; i < result_rank; i++) {
     int64_t current_bound;
     bool dim_dynamic = false;
+    DExpr expression = DExpr::Unknown(80);
     bool is_window_index =
         absl::c_binary_search(gather_dim_numbers.offset_dims(), i);
     if (is_window_index) {
@@ -4353,6 +4528,9 @@ static absl::Status ValidateGatherDimensionNumbers(
       if (slice_sizes[offset_dims_seen] ==
           input_shape.dimensions(offset_dims_seen)) {
         dim_dynamic = input_shape.is_dynamic_dimension(offset_dims_seen);
+        expression = input_shape.expressions(offset_dims_seen);
+      } else {
+        expression = DExpr::Const(slice_sizes[offset_dims_seen]);
       }
       current_bound = slice_sizes[offset_dims_seen++];
     } else {
@@ -4362,15 +4540,17 @@ static absl::Status ValidateGatherDimensionNumbers(
       // Forward dynamic dimensions from indices.
       dim_dynamic =
           expanded_start_indices_shape_dynamic_dimensions[gather_dims_seen];
-
+      expression = expanded_start_indices_shape_expressions[gather_dims_seen];
       current_bound = expanded_start_indices_shape[gather_dims_seen++];
     }
     output_dim_is_dynamic.push_back(dim_dynamic);
+    output_expressions.push_back(expression);
     output_dim_bounds.push_back(current_bound);
   }
 
-  return ShapeUtil::MakeShape(input_shape.element_type(), output_dim_bounds,
-                              output_dim_is_dynamic);
+  auto s = ShapeUtil::MakeShape(input_shape.element_type(), output_dim_bounds,
+                                output_dim_is_dynamic, output_expressions);
+  return s;
 }
 
 namespace {
diff --git a/third_party/xla/xla/service/shape_inference.h b/third_party/xla/xla/service/shape_inference.h
index 15818f01cdfa7e..cb0ee9a20f83c6 100644
--- a/third_party/xla/xla/service/shape_inference.h
+++ b/third_party/xla/xla/service/shape_inference.h
@@ -245,13 +245,17 @@ class ShapeInference {
   // e.g. slice f32[32x32] 0:16 0:16 -> f32[16x16]
   static absl::StatusOr<Shape> InferSliceShape(
       const Shape& arg, absl::Span<const int64_t> starts,
-      absl::Span<const int64_t> limits, absl::Span<const int64_t> strides);
+      absl::Span<const int64_t> limits, absl::Span<const int64_t> strides,
+      absl::Span<const DExpr> start_exprs = {},
+      absl::Span<const DExpr> limit_exprs = {});
 
   // Infers the shape produced by a dynamic slice operation of size specified
   // in 'slice_sizes', with dynamic start indices shape 'start_indices_shape'.
   static absl::StatusOr<Shape> InferDynamicSliceShape(
       const Shape& operand_shape, absl::Span<const Shape> start_index_shapes,
-      absl::Span<const int64_t> slice_sizes, bool allow_scalar_indices = true);
+      absl::Span<const int64_t> slice_sizes,
+      absl::Span<const DExpr> slice_exprs = {},
+      bool allow_scalar_indices = true);
 
   // Infers the shape produced by a dynamic update slice operation based
   // on the shape of operand and update.
@@ -283,7 +287,8 @@ class ShapeInference {
 
   // Infers the shape produced by a broadcast operation.
   static absl::StatusOr<Shape> InferBroadcastShape(
-      const Shape& operand, absl::Span<const int64_t> broadcast_sizes);
+      const Shape& operand, absl::Span<const int64_t> broadcast_sizes,
+      absl::Span<const DExpr> broadcast_exprs = {});
 
   // Checks whether the given parameters can form a broadcast. Returns the same
   // output_shape if it's legal.
@@ -295,7 +300,7 @@ class ShapeInference {
   // its operand and the new dimension sizes specified.
   static absl::StatusOr<Shape> InferReshapeShape(
       const Shape& operand, absl::Span<const int64_t> dimensions,
-      int64_t inferred_dimension);
+      int64_t inferred_dimension, absl::Span<const DExpr> expressions);
 
   // Infers the shape produced by a dynamic reshape operation from the element
   // type of its operand and the new dimension sizes specified. The result shape
@@ -304,7 +309,8 @@ class ShapeInference {
   static absl::StatusOr<Shape> InferDynamicReshapeShape(
       const Shape& operand, absl::Span<const Shape* const> dim_size_shapes,
       absl::Span<const int64_t> new_size_bounds,
-      const std::vector<bool>& dims_are_dynamic);
+      const std::vector<bool>& dims_are_dynamic,
+      absl::Span<const DExpr> expressions);
 
   // Infers the shape produced by a transpose operation from the element type of
   // its operand and its dimensions field.
diff --git a/third_party/xla/xla/service/shape_inference_test.cc b/third_party/xla/xla/service/shape_inference_test.cc
index 87841d476712d7..0b6e4bf00650b2 100644
--- a/third_party/xla/xla/service/shape_inference_test.cc
+++ b/third_party/xla/xla/service/shape_inference_test.cc
@@ -205,6 +205,47 @@ TEST_F(ShapeInferenceTest, SelectArrayPredBetweenArrays) {
   ASSERT_TRUE(ShapeUtil::Equal(matrix_64_48_, *inferred_shape));
 }
 
+TEST_F(ShapeInferenceTest, SelectPreservesExpressionsFromOperands) {
+  const Shape pred = ShapeUtil::MakeShape(PRED, {8, 5});
+  const Shape on_true = ShapeUtil::MakeShape(
+      F32, {8, 5}, std::vector<DExpr>{DExpr::Var(33), DExpr::Const(5)});
+  const Shape on_false = ShapeUtil::MakeShape(
+      F32, {8, 5}, std::vector<DExpr>{DExpr::Var(33), DExpr::Const(5)});
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred, on_true,
+                                          on_false));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(33)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(5)));
+}
+
+TEST_F(ShapeInferenceTest, SelectRejectsMismatchedOperandExpressions) {
+  const Shape pred = ShapeUtil::MakeShape(PRED, {});
+  const Shape on_true = ShapeUtil::MakeShape(
+      F32, {8, 5}, std::vector<DExpr>{DExpr::Var(33), DExpr::Const(5)});
+  const Shape on_false = ShapeUtil::MakeShape(
+      F32, {8, 5}, std::vector<DExpr>{DExpr::Var(34), DExpr::Const(5)});
+  const absl::StatusOr<Shape> inferred_shape =
+      ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred, on_true,
+                                          on_false);
+  ASSERT_FALSE(inferred_shape.ok());
+  EXPECT_THAT(inferred_shape.status().message(),
+              HasSubstr("mismatched expressions in dimension 0"));
+}
+
+TEST_F(ShapeInferenceTest, SelectRejectsMissingDynamicOperandExpression) {
+  const Shape pred = ShapeUtil::MakeShape(PRED, {});
+  const Shape on_true = ShapeUtil::MakeShape(
+      F32, {8, 5}, std::vector<DExpr>{DExpr::Var(33), DExpr::Const(5)});
+  const Shape on_false = ShapeUtil::MakeShape(F32, {8, 5});
+  const absl::StatusOr<Shape> inferred_shape =
+      ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred, on_true,
+                                          on_false);
+  ASSERT_FALSE(inferred_shape.ok());
+  EXPECT_THAT(inferred_shape.status().message(),
+              HasSubstr("mismatched expressions in dimension 0"));
+}
+
 TEST_F(ShapeInferenceTest, SelectBadShapes) {
   const absl::StatusOr<Shape> inferred_shape_error1 =
       ShapeInference::InferTernaryOpShape(HloOpcode::kSelect, pred_,
@@ -448,6 +489,63 @@ TEST_F(ShapeInferenceTest, ReduceWindowInHalf) {
       ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {4, 4}), *inferred_shape));
 }
 
+TEST_F(ShapeInferenceTest, ReduceWindowBuildsWindowedExpressions) {
+  const Shape matrix_shape = ShapeUtil::MakeShape(
+      F32, {8, 8}, std::vector<DExpr>{DExpr::Var(1), DExpr::Var(2)});
+  Window window;
+  WindowDimension dim;
+  dim.set_size(2);
+  dim.set_stride(2);
+  dim.set_padding_low(0);
+  dim.set_padding_high(0);
+  dim.set_window_dilation(1);
+  dim.set_base_dilation(1);
+  *window.add_dimensions() = dim;
+  *window.add_dimensions() = dim;
+  const Shape init_value_shape = ShapeUtil::MakeShape(F32, {});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReduceWindowShape(matrix_shape, init_value_shape,
+                                             window));
+
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape, ShapeUtil::MakeShape(
+                          F32, {4, 4},
+                          std::vector<DExpr>{((DExpr::Var(1) - 2) / 2) + 1,
+                                             ((DExpr::Var(2) - 2) / 2) + 1})));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             ((DExpr::Var(1) - 2) / 2) + 1));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1),
+                             ((DExpr::Var(2) - 2) / 2) + 1));
+}
+
+TEST_F(ShapeInferenceTest, ReduceWindowBuildsPaddedStridedExpressions) {
+  const Shape vector_shape =
+      ShapeUtil::MakeShape(F32, {11}, std::vector<DExpr>{DExpr::Var(3)});
+  Window window;
+  WindowDimension dim;
+  dim.set_size(3);
+  dim.set_stride(2);
+  dim.set_padding_low(1);
+  dim.set_padding_high(1);
+  dim.set_window_dilation(1);
+  dim.set_base_dilation(1);
+  *window.add_dimensions() = dim;
+  const Shape init_value_shape = ShapeUtil::MakeShape(F32, {});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReduceWindowShape(vector_shape, init_value_shape,
+                                             window));
+
+  const DExpr expected = (((DExpr::Var(3) - 1) / 2) + 1).simplify();
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape,
+      ShapeUtil::MakeShape(F32, {6}, std::vector<DExpr>{expected})));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), expected));
+}
+
 TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterProperShapes) {
   const absl::StatusOr<Shape> inferred_shape_ok =
       ShapeInference::InferSelectAndScatterShape(
@@ -721,6 +819,58 @@ TEST_F(ShapeInferenceTest, ConvolveWithBaseDilation) {
                                *inferred_shape));
 }
 
+TEST_F(ShapeInferenceTest, ConvolveBuildsBatchAndSpatialExpressions) {
+  ConvolutionDimensionNumbers dnums;
+  const Shape lhs_shape = ShapeUtil::MakeShape(
+      F32, {5, 11, 13, 3},
+      std::vector<DExpr>{DExpr::Var(4), DExpr::Var(5), DExpr::Var(6),
+                         DExpr::Const(3)});
+  dnums.set_input_batch_dimension(0);
+  dnums.set_output_batch_dimension(0);
+  dnums.add_input_spatial_dimensions(1);
+  dnums.add_output_spatial_dimensions(1);
+  dnums.add_input_spatial_dimensions(2);
+  dnums.add_output_spatial_dimensions(2);
+  dnums.set_input_feature_dimension(3);
+  dnums.set_output_feature_dimension(3);
+
+  const Shape rhs_shape = ShapeUtil::MakeShape(F32, {3, 3, 3, 7});
+  dnums.add_kernel_spatial_dimensions(0);
+  dnums.add_kernel_spatial_dimensions(1);
+  dnums.set_kernel_input_feature_dimension(2);
+  dnums.set_kernel_output_feature_dimension(3);
+
+  Window window;
+  auto* dim0 = window.add_dimensions();
+  dim0->set_size(3);
+  dim0->set_stride(2);
+  dim0->set_padding_low(1);
+  dim0->set_padding_high(1);
+  dim0->set_window_dilation(1);
+  dim0->set_base_dilation(1);
+  auto* dim1 = window.add_dimensions();
+  *dim1 = *dim0;
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferConvolveShape(
+          lhs_shape, rhs_shape, /*feature_group_count=*/1,
+          /*batch_group_count=*/1, window, dnums,
+          /*preferred_element_type=*/std::nullopt));
+
+  const DExpr expected_h = (((DExpr::Var(5) - 1) / 2) + 1).simplify();
+  const DExpr expected_w = (((DExpr::Var(6) - 1) / 2) + 1).simplify();
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape,
+      ShapeUtil::MakeShape(F32, {5, 6, 7, 7},
+                           std::vector<DExpr>{DExpr::Var(4), expected_h,
+                                              expected_w, DExpr::Const(7)})));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(4)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), expected_h));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), expected_w));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(3), DExpr::Const(7)));
+}
+
 TEST_F(ShapeInferenceTest, ConvolveDimensionNumbersOverlapError) {
   // Dimension order for this test: batch, feature, x0, x1
   const Shape lhs_shape = ShapeUtil::MakeShape(F32, {10, 11, 3, 4});
@@ -1028,7 +1178,8 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestFftRanks) {
 
 TEST_F(ShapeInferenceTest, InferFftShapeTestFftRanksBounded) {
   FftType type = FftType::FFT;
-  const Shape shape = ShapeUtil::MakeShape(C64, {16, 8}, {false, true});
+  const Shape shape =
+      ShapeUtil::MakeShape(C64, {16, 8}, {false, true}, {});
   fft::Fail(shape, type, {}, fft::unsupported_rank);
   fft::Pass(shape, type, {8}, shape);
   fft::Pass(shape, type, {16, 8}, shape);
@@ -1056,7 +1207,8 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestIfftRanks) {
 
 TEST_F(ShapeInferenceTest, InferFftShapeTestIfftRanksBounded) {
   FftType type = FftType::IFFT;
-  const Shape shape = ShapeUtil::MakeShape(C64, {16, 8}, {false, true});
+  const Shape shape =
+      ShapeUtil::MakeShape(C64, {16, 8}, {false, true}, {});
   fft::Fail(shape, type, {}, fft::unsupported_rank);
   fft::Pass(shape, type, {8}, shape);
   fft::Pass(shape, type, {16, 8}, shape);
@@ -1103,9 +1255,9 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestRfftDimensions) {
   fft::Pass(odd_shape_in, type, {16, 9}, shape_out);
 
   const Shape bounded_shape_in =
-      ShapeUtil::MakeShape(F32, {16, 8}, {false, true});
+      ShapeUtil::MakeShape(F32, {16, 8}, {false, true}, {});
   const Shape bounded_shape_out =
-      ShapeUtil::MakeShape(C64, {16, 5}, {false, true});
+      ShapeUtil::MakeShape(C64, {16, 5}, {false, true}, {});
   fft::Pass(bounded_shape_in, type, {16, 8}, bounded_shape_out);
 }
 
@@ -1147,9 +1299,9 @@ TEST_F(ShapeInferenceTest, InferFftShapeTestIrfftDimensions) {
   fft::Pass(shape, type, {16, 9}, odd_shape_out);
 
   const Shape bounded_shape_in =
-      ShapeUtil::MakeShape(C64, {16, 5}, {false, true});
+      ShapeUtil::MakeShape(C64, {16, 5}, {false, true}, {});
   const Shape bounded_shape_out =
-      ShapeUtil::MakeShape(F32, {16, 9}, {false, true});
+      ShapeUtil::MakeShape(F32, {16, 9}, {false, true}, {});
   fft::Pass(bounded_shape_in, type, {16, 9}, bounded_shape_out);
 }
 
@@ -1283,6 +1435,18 @@ TEST_F(ShapeInferenceTest, MapWithDifferentInputTypes) {
   EXPECT_TRUE(ShapeUtil::Equal(expected, *inferred_shape));
 }
 
+TEST_F(ShapeInferenceTest, MapPreservesExpressions) {
+  const Shape arg = ShapeUtil::MakeShape(
+      F32, {20, 7}, std::vector<bool>{true, false},
+      std::vector<DExpr>{DExpr::Var(11), DExpr::Const(7)});
+  ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_}, f32_);
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferMapShape({&arg}, to_apply,
+                                                        {0, 1}));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(11)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(7)));
+}
+
 TEST_F(ReduceShapeInferenceTest, ReduceVectorToScalar) {
   ExpectInferredReduceShape(f32_, ShapeUtil::MakeShape(F32, {128}),
                             /*dimensions_to_reduce=*/{0});
@@ -1367,6 +1531,58 @@ TEST_F(ReduceShapeInferenceTest, ReduceWindowMultiOutput) {
       *inferred_shape));
 }
 
+TEST_F(ReduceShapeInferenceTest,
+       ReduceWindowPreservesDynamicStridedBoundExpression) {
+  Shape operand = ShapeUtil::MakeShape(
+      F32, {101}, std::vector<bool>{true}, {DExpr::Var(1)});
+  Window window;
+  WindowDimension* dimension = window.add_dimensions();
+  dimension->set_size(3);
+  dimension->set_stride(2);
+  dimension->set_padding_low(1);
+  dimension->set_padding_high(1);
+  dimension->set_base_dilation(1);
+  dimension->set_window_dilation(1);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      Shape inferred,
+      ShapeInference::InferReduceWindowShape(operand, f32_, window));
+  EXPECT_EQ(51, inferred.dimensions(0));
+  EXPECT_TRUE(inferred.expressions(0) ==
+              DExpr::Max((DExpr::Var(1) + 1) / 2, DExpr::Const(0)));
+
+  DExpr runtime_expression =
+      inferred.expressions(0).substitute(1, DExpr::Const(100)).simplify();
+  ASSERT_TRUE(runtime_expression->is_constant());
+  EXPECT_EQ(50, runtime_expression->get_val());
+}
+
+TEST_F(ReduceShapeInferenceTest,
+       ReduceWindowClampsDynamicStridedBoundAtZero) {
+  Shape operand = ShapeUtil::MakeShape(
+      F32, {101}, std::vector<bool>{true}, {DExpr::Var(2)});
+  Window window;
+  WindowDimension* dimension = window.add_dimensions();
+  dimension->set_size(5);
+  dimension->set_stride(1);
+  dimension->set_padding_low(0);
+  dimension->set_padding_high(0);
+  dimension->set_base_dilation(1);
+  dimension->set_window_dilation(1);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      Shape inferred,
+      ShapeInference::InferReduceWindowShape(operand, f32_, window));
+  EXPECT_EQ(97, inferred.dimensions(0));
+  EXPECT_TRUE(inferred.expressions(0) ==
+              DExpr::Max(DExpr::Var(2) - 4, DExpr::Const(0)));
+
+  DExpr runtime_expression =
+      inferred.expressions(0).substitute(2, DExpr::Const(2)).simplify();
+  ASSERT_TRUE(runtime_expression->is_constant());
+  EXPECT_EQ(0, runtime_expression->get_val());
+}
+
 TEST_F(ReduceShapeInferenceTest, ErrorMultiOutputBadReducerInput1) {
   const Shape f32_arg_shape = ShapeUtil::MakeShape(F32, {5, 3});
   const Shape s32_arg_shape = ShapeUtil::MakeShape(S32, {5, 3});
@@ -1520,12 +1736,31 @@ TEST_F(ShapeInferenceTest, InferSliceShapeRank2) {
 }
 
 TEST_F(ShapeInferenceTest, InferSliceWithDynamicDimensions) {
-  const Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64}, {true, true});
+  const Shape matrix_shape =
+      ShapeUtil::MakeShape(F32, {128, 64}, {true, true}, {});
   const absl::StatusOr<Shape> inferred_shape =
       ShapeInference::InferSliceShape(matrix_shape, {32, 0}, {33, 64}, {1, 1});
   ASSERT_IS_OK(inferred_shape.status());
   ASSERT_TRUE(ShapeUtil::Equal(
-      ShapeUtil::MakeShape(F32, {1, 64}, {false, true}), *inferred_shape));
+      ShapeUtil::MakeShape(F32, {1, 64}, {false, true}, {}),
+      *inferred_shape));
+}
+
+TEST_F(ShapeInferenceTest, InferSliceBuildsExpressionFromSymbolicBounds) {
+  const Shape vector_shape =
+      ShapeUtil::MakeShape(F32, {16}, std::vector<DExpr>{DExpr::Var(1)});
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferSliceShape(
+          vector_shape, /*starts=*/{3}, /*limits=*/{8}, /*strides=*/{1},
+          /*start_exprs=*/{DExpr::Const(3)},
+          /*limit_exprs=*/{DExpr::Var(2)}));
+
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape, ShapeUtil::MakeShape(
+                          F32, {5}, std::vector<DExpr>{DExpr::Var(2) - 3})));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(2) - 3));
 }
 
 TEST_F(ShapeInferenceTest, InferSliceShapeRank2WithStrides) {
@@ -1583,6 +1818,26 @@ TEST_F(ShapeInferenceTest, InferConstIndexShape) {
   ASSERT_TRUE(ShapeUtil::Equal(s32_, *inferred1_status));
 }
 
+TEST_F(ShapeInferenceTest, InferConstIndexPreservesExpressions) {
+  const Shape lhs = ShapeUtil::MakeShape(
+      F32, {8, 5}, std::vector<DExpr>{DExpr::Var(24), DExpr::Const(5)});
+  const Shape rhs = ShapeUtil::MakeShape(
+      S32, {3, 7}, std::vector<DExpr>{DExpr::Const(3), DExpr::Var(25)});
+  const Shape tuple_shape = ShapeUtil::MakeTupleShape({lhs, rhs});
+
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred0,
+                          ShapeInference::InferGetTupleElementShape(
+                              tuple_shape, /*index=*/0));
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred1,
+                          ShapeInference::InferGetTupleElementShape(
+                              tuple_shape, /*index=*/1));
+
+  EXPECT_TRUE(DynExpr::equal(inferred0.expressions(0), DExpr::Var(24)));
+  EXPECT_TRUE(DynExpr::equal(inferred0.expressions(1), DExpr::Const(5)));
+  EXPECT_TRUE(DynExpr::equal(inferred1.expressions(0), DExpr::Const(3)));
+  EXPECT_TRUE(DynExpr::equal(inferred1.expressions(1), DExpr::Var(25)));
+}
+
 TEST_F(ShapeInferenceTest, InferTupleElementShapeOutOfBound) {
   const Shape tuple_shape = ShapeUtil::MakeTupleShape({f32_, s32_});
   const absl::StatusOr<Shape> inferredNegative_status =
@@ -1622,11 +1877,13 @@ TEST_F(ShapeInferenceTest, InferReshapeDegenerateCombine) {
   // [<=1]
   //
   // Both output dimension can be dynamic, use inferred_dimension to tie-break.
-  const Shape operand = ShapeUtil::MakeShape(F32, {1, 1}, {false, true});
+  const Shape operand = ShapeUtil::MakeShape(F32, {1, 1}, {false, true}, {});
   const auto status =
       ShapeInference::InferReshapeShape(operand, {1},
-                                        /*inferred_dimension=*/-1);
-  ASSERT_EQ(ShapeUtil::MakeShape(F32, {1}, {true}), *status);
+                                        /*inferred_dimension=*/-1,
+                                        /*expressions=*/{});
+  ASSERT_EQ(ShapeUtil::MakeShape(F32, {1}, {true}, {}),
+            *status);
 }
 
 TEST_F(ShapeInferenceTest, InferReshapeSplit) {
@@ -1635,48 +1892,218 @@ TEST_F(ShapeInferenceTest, InferReshapeSplit) {
   // [1, 10]
   //
   // Both output dimension can be dynamic, use inferred_dimension to tie-break.
-  const Shape operand = ShapeUtil::MakeShape(F32, {10}, {true});
+  const Shape operand = ShapeUtil::MakeShape(F32, {10}, {true}, {});
   const auto status =
       ShapeInference::InferReshapeShape(operand, {1, 10},
-                                        /*inferred_dimension=*/0);
-  ASSERT_EQ(ShapeUtil::MakeShape(F32, {1, 10}, {true, false}), *status);
+                                        /*inferred_dimension=*/0,
+                                        /*expressions=*/{});
+  ASSERT_EQ(
+      ShapeUtil::MakeShape(F32, {1, 10}, {true, false}, {}),
+      *status);
 }
 
 TEST_F(ShapeInferenceTest, InferReshapeCombine) {
   // [6, <=10]
   //   | reshape
   // [<=60]
-  const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true});
+  const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true}, {});
   const auto status =
       ShapeInference::InferReshapeShape(operand, {60},
-                                        /*inferred_dimension=*/-11);
-  ASSERT_EQ(ShapeUtil::MakeShape(F32, {60}, {true}), *status);
+                                        /*inferred_dimension=*/-11,
+                                        /*expressions=*/{});
+  ASSERT_EQ(ShapeUtil::MakeShape(F32, {60}, {true}, {}),
+            *status);
 }
 
 TEST_F(ShapeInferenceTest, UnchangedDimension) {
   // [6, <=10]
   //   | reshape
   // [2, 3, <=10]
-  const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true});
+  const Shape operand = ShapeUtil::MakeShape(F32, {6, 10}, {false, true}, {});
   const auto status =
       ShapeInference::InferReshapeShape(operand, {2, 3, 10},
-                                        /*inferred_dimension=*/-11);
-  ASSERT_EQ(ShapeUtil::MakeShape(F32, {2, 3, 10}, {false, false, true}),
+                                        /*inferred_dimension=*/-11,
+                                        /*expressions=*/{});
+  ASSERT_EQ(ShapeUtil::MakeShape(F32, {2, 3, 10}, {false, false, true}, {}),
             *status);
 }
 
+TEST_F(ShapeInferenceTest, ReshapePreservesProvidedExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {6, 10}, std::vector<DExpr>{DExpr::Const(6), DExpr::Var(1)});
+  const Shape expected = ShapeUtil::MakeShape(
+      F32, {2, 3, 10},
+      std::vector<DExpr>{DExpr::Const(2), DExpr::Const(3), DExpr::Var(1)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReshapeShape(
+          operand, expected.dimensions(),
+          /*inferred_dimension=*/-1, expected.expressions()));
+
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Const(2)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Var(1)));
+}
+
+TEST_F(ShapeInferenceTest, ReshapeCombinesLeadingSymbolicWithStaticFactor) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {5, 16, 32},
+      std::vector<DExpr>{DExpr::Var(1), DExpr::Const(16), DExpr::Const(32)});
+  const Shape expected = ShapeUtil::MakeShape(
+      F32, {80, 32},
+      std::vector<DExpr>{16 * DExpr::Var(1), DExpr::Const(32)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReshapeShape(
+          operand, expected.dimensions(),
+          /*inferred_dimension=*/-1, expected.expressions()));
+
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             16 * DExpr::Var(1)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(32)));
+}
+
+TEST_F(ShapeInferenceTest, ReshapeCollapsesTwoStaticDimsIntoSymbolicExtent) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {5, 4, 8},
+      std::vector<DExpr>{DExpr::Var(1), DExpr::Const(4), DExpr::Const(8)});
+  const Shape expected =
+      ShapeUtil::MakeShape(F32, {160}, std::vector<DExpr>{32 * DExpr::Var(1)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReshapeShape(
+          operand, expected.dimensions(),
+          /*inferred_dimension=*/-1, expected.expressions()));
+
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             32 * DExpr::Var(1)));
+}
+
+TEST_F(ShapeInferenceTest, ReshapeSplitsSymbolicExtentByStaticFactor) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {80, 8}, std::vector<DExpr>{DExpr::Var(1), DExpr::Const(8)});
+  const Shape expected = ShapeUtil::MakeShape(
+      F32, {40, 16},
+      std::vector<DExpr>{DExpr::Var(1) / 2, DExpr::Const(16)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReshapeShape(
+          operand, expected.dimensions(),
+          /*inferred_dimension=*/-1, expected.expressions()));
+
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             DExpr::Var(1) / 2));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(16)));
+}
+
+TEST_F(ShapeInferenceTest, ReshapeSplitsAndCollapsesSymbolicExtent) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {20, 8, 4},
+      std::vector<DExpr>{DExpr::Var(1), DExpr::Const(8), DExpr::Const(4)});
+  const Shape expected = ShapeUtil::MakeShape(
+      F32, {10, 64},
+      std::vector<DExpr>{DExpr::Var(1) / 2, DExpr::Const(64)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReshapeShape(
+          operand, expected.dimensions(),
+          /*inferred_dimension=*/-1, expected.expressions()));
+
+  EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             DExpr::Var(1) / 2));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(64)));
+}
+
+TEST_F(ShapeInferenceTest, ReshapeRejectsIncorrectCollapsedExpression) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {5, 16, 32},
+      std::vector<DExpr>{DExpr::Var(1), DExpr::Const(16), DExpr::Const(32)});
+  const Shape incorrect = ShapeUtil::MakeShape(
+      F32, {80, 32},
+      std::vector<DExpr>{8 * DExpr::Var(1), DExpr::Const(32)});
+
+  const absl::StatusOr<Shape> status = ShapeInference::InferReshapeShape(
+      operand, incorrect.dimensions(),
+      /*inferred_dimension=*/-1, incorrect.expressions());
+
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.status().message(),
+              HasSubstr("Reshape operation has mismatched symbolic element "
+                        "counts"));
+}
+
+TEST_F(ShapeInferenceTest,
+       ReshapeRejectsIncorrectSplitAndCollapseExpression) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {20, 8, 4},
+      std::vector<DExpr>{DExpr::Var(1), DExpr::Const(8), DExpr::Const(4)});
+  const Shape incorrect = ShapeUtil::MakeShape(
+      F32, {10, 64},
+      std::vector<DExpr>{DExpr::Var(1), DExpr::Const(64)});
+
+  const absl::StatusOr<Shape> status = ShapeInference::InferReshapeShape(
+      operand, incorrect.dimensions(),
+      /*inferred_dimension=*/-1, incorrect.expressions());
+
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.status().message(),
+              HasSubstr("Reshape operation has mismatched symbolic element "
+                        "counts"));
+}
+
+TEST_F(ShapeInferenceTest, ReshapeWithSymbolicOperandRequiresExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {10, 6}, std::vector<DExpr>{DExpr::Var(1), DExpr::Const(6)});
+
+  const absl::StatusOr<Shape> status =
+      ShapeInference::InferReshapeShape(operand, {2, 5, 6},
+                                        /*inferred_dimension=*/-1,
+                                        /*expressions=*/{});
+
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.status().message(),
+              HasSubstr("Expressions is empty but operand is dynamic"));
+}
+
 TEST_F(ShapeInferenceTest, InferDynamicBroadcast) {
   // CHECK:
   // %broadcast = s32[15,<=15]{1,0} broadcast(s32[<=15]{0}), dimensions={1}
 
-  const Shape operand_shape = ShapeUtil::MakeShape(F32, {15}, {true});
+  const Shape operand_shape = ShapeUtil::MakeShape(F32, {15}, {true}, {});
   const absl::StatusOr<Shape> inferred_shape =
       ShapeInference::InferBroadcastShape(operand_shape, {15});
   ASSERT_IS_OK(inferred_shape.status());
-  ASSERT_EQ(ShapeUtil::MakeShape(F32, {15, 15}, {false, true}),
+  ASSERT_EQ(ShapeUtil::MakeShape(F32, {15, 15}, {false, true}, {}),
             *inferred_shape);
 }
 
+TEST_F(ShapeInferenceTest, BroadcastInDimPreservesMappedExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {2, 4}, std::vector<bool>{true, true},
+      {DExpr::Var(1), DExpr::Var(2)});
+  const Shape output = ShapeUtil::MakeShape(
+      F32, {2, 3, 4}, std::vector<bool>{true, false, true},
+      {DExpr::Var(1), DExpr::Const(3), DExpr::Var(2)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferBroadcastShape(operand, output,
+                                          /*broadcast_dimensions=*/{0, 2}));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(1)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Var(2)));
+}
+
 TEST_F(ShapeInferenceTest, BroadcastScalar) {
   for (auto element_type : {F32, U32, S8}) {
     const Shape scalar_shape = ShapeUtil::MakeShape(element_type, {});
@@ -2172,6 +2599,27 @@ TEST_F(ShapeInferenceTest, SparseDotMetadata) {
       ShapeUtil::Equal(inferred_shape, ShapeUtil::MakeShape(U16, {5, 10, 2})));
 }
 
+TEST_F(ShapeInferenceTest, SparseDotMetadataPreservesNonSparseExpressions) {
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_batch_dimensions(0);
+  dot_dnums.add_lhs_contracting_dimensions(2);
+  SparsityDescriptor sparsity_descriptor;
+  sparsity_descriptor.set_type(SparsityType::SPARSITY_STRUCTURED_N_M);
+  sparsity_descriptor.set_n(2);
+  sparsity_descriptor.set_m(4);
+  sparsity_descriptor.set_index(0);
+  sparsity_descriptor.set_dimension(2);
+
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {5, 10, 16},
+      std::vector<DExpr>{DExpr::Var(16), DExpr::Const(10), DExpr::Const(16)});
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferSparseDotMetadataShape(
+                              operand, dot_dnums, sparsity_descriptor));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(16)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(10)));
+}
+
 // <ragged-dot> mode 1 : [m,k], [g,k,n], [g] -> [m,n]
 TEST_F(ShapeInferenceTest, RaggedDotRaggedNonContracting) {
   const Shape lhs_shape = ShapeUtil::MakeShape(F32, {11, 5});
@@ -2221,6 +2669,31 @@ TEST_F(ShapeInferenceTest, RaggedDotRaggedContracting) {
       << " expected: " << ShapeUtil::HumanString(output_shape);
 }
 
+TEST_F(ShapeInferenceTest, RaggedDotPreservesGroupExpression) {
+  const Shape lhs_shape = ShapeUtil::MakeShape(
+      F32, {11, 5}, std::vector<DExpr>{DExpr::Const(11), DExpr::Const(5)});
+  const Shape rhs_shape = ShapeUtil::MakeShape(
+      F32, {5, 7}, std::vector<DExpr>{DExpr::Const(5), DExpr::Const(7)});
+  const Shape group_sizes_shape =
+      ShapeUtil::MakeShape(U32, {3}, std::vector<DExpr>{DExpr::Var(17)});
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(1);
+  dot_dnums.add_rhs_contracting_dimensions(0);
+  RaggedDotDimensionNumbers ragged_dot_dnums;
+  *ragged_dot_dnums.mutable_dot_dimension_numbers() = dot_dnums;
+  ragged_dot_dnums.add_lhs_ragged_dimensions(1);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferRaggedDotOpShape(
+          lhs_shape, rhs_shape, group_sizes_shape, ragged_dot_dnums,
+          /*preferred_element_type=*/std::nullopt));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(17)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(11)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(7)));
+}
+
 // <ragged-dot> mode 3 : [b,m,k], [b,k,n], [g] -> [b,m,n]
 TEST_F(ShapeInferenceTest, RaggedDotRaggedBatch) {
   const Shape lhs_shape = ShapeUtil::MakeShape(F32, {3, 11, 5});
@@ -2727,6 +3200,54 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastMatrixVector) {
   ASSERT_FALSE(inferred_shape_mismatch.ok());
 }
 
+TEST_F(ShapeInferenceTest, InDimBroadcastPreservesMappedExpressions) {
+  const Shape smaller = ShapeUtil::MakeShape(
+      F32, {2, 3},
+      std::vector<DExpr>{DExpr::Var(18), DExpr::Const(3)});
+  const Shape larger = ShapeUtil::MakeShape(
+      F32, {2, 4, 3},
+      std::vector<DExpr>{DExpr::Const(2), DExpr::Const(4), DExpr::Const(3)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, larger, smaller,
+                                         {0, 2}));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(18)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(4)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(3)));
+}
+
+TEST_F(ShapeInferenceTest, DegenerateBroadcastUsesNonUnitExpressions) {
+  const Shape lhs = ShapeUtil::MakeShape(
+      F32, {1, 3},
+      std::vector<DExpr>{DExpr::Const(1), DExpr::Const(3)});
+  const Shape rhs = ShapeUtil::MakeShape(
+      F32, {5, 3},
+      std::vector<DExpr>{DExpr::Var(19), DExpr::Const(3)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, lhs, rhs, {}));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(19)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3)));
+}
+
+TEST_F(ShapeInferenceTest, ElementwiseBinaryBroadcastPreservesExpressions) {
+  const Shape lhs = ShapeUtil::MakeShape(
+      F32, {2, 4, 3},
+      std::vector<DExpr>{DExpr::Const(2), DExpr::Const(4), DExpr::Const(3)});
+  const Shape rhs = ShapeUtil::MakeShape(
+      F32, {2, 3},
+      std::vector<DExpr>{DExpr::Var(20), DExpr::Const(3)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, lhs, rhs, {0, 2}));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(20)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(4)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(3)));
+}
+
 TEST_F(ShapeInferenceTest, BinOpBroadcastCubeMatrix) {
   // Test variations of broadcasting a matrix for a binary add with a cube.
   const Shape cube = ShapeUtil::MakeShape(F32, {16, 8, 4});
@@ -2822,6 +3343,17 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastBadDimension) {
               HasSubstr("dimensions order is wrong"));
 }
 
+TEST_F(ShapeInferenceTest, BinOpPreservesBroadcastedExpressionSameRank) {
+  Shape lhs = ShapeUtil::MakeShape(F32, {1});
+  Shape rhs =
+      ShapeUtil::MakeShape(F32, {1}, std::vector<bool>{true}, {DExpr::Var(1)});
+
+  const absl::StatusOr<Shape> inferred_shape =
+      ShapeInference::InferBinaryOpShape(HloOpcode::kAdd, lhs, rhs, {0});
+  ASSERT_IS_OK(inferred_shape.status());
+  EXPECT_EQ(inferred_shape->expressions(0), DExpr::Var(1));
+}
+
 // Tests for the while instruction with proper shapes.
 TEST_F(ShapeInferenceTest, WhileWithCorrectShapes) {
   const Shape result_shape = ShapeUtil::MakeTupleShape({s32_, vector_32_});
@@ -2875,18 +3407,38 @@ TEST_F(ShapeInferenceTest, WhileWithBadShapes) {
 // Tests for the concatenate instruction with dynamic shapes.
 TEST_F(ShapeInferenceTest, ConcatenateWithDynamicShapes) {
   const auto dynamic_shape_1 =
-      ShapeUtil::MakeShape(F32, {32, 160, 10}, {true, false, false});
+      ShapeUtil::MakeShape(F32, {32, 160, 10}, {true, false, false}, {});
   const auto dynamic_shape_2 =
-      ShapeUtil::MakeShape(F32, {32, 160, 10}, {false, true, false});
+      ShapeUtil::MakeShape(F32, {32, 160, 10}, {false, true, false}, {});
   const absl::StatusOr<Shape> inferred_shape =
       ShapeInference::InferConcatOpShape({&dynamic_shape_1, &dynamic_shape_2},
                                          /*dimension=*/0);
   ASSERT_IS_OK(inferred_shape.status());
   ASSERT_TRUE(ShapeUtil::Equal(
-      ShapeUtil::MakeShape(F32, {64, 160, 10}, {true, true, false}),
+      ShapeUtil::MakeShape(F32, {64, 160, 10}, {true, true, false}, {}),
       *inferred_shape));
 }
 
+TEST_F(ShapeInferenceTest, ConcatenateAddsConcatDimensionExpressions) {
+  const Shape lhs = ShapeUtil::MakeShape(
+      F32, {2, 5}, std::vector<DExpr>{DExpr::Var(1), DExpr::Const(5)});
+  const Shape rhs = ShapeUtil::MakeShape(
+      F32, {3, 5}, std::vector<DExpr>{DExpr::Var(2), DExpr::Const(5)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferConcatOpShape({&lhs, &rhs}, /*dimension=*/0));
+
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape,
+      ShapeUtil::MakeShape(F32, {5, 5},
+                           std::vector<DExpr>{DExpr::Var(1) + DExpr::Var(2),
+                                              DExpr::Const(5)})));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             DExpr::Var(1) + DExpr::Var(2)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(5)));
+}
+
 // Tests for the concatenate instruction with proper shapes.
 TEST_F(ShapeInferenceTest, ConcatenateWithCorrectShapes) {
   const absl::StatusOr<Shape> inferred_shape_1 =
@@ -2988,6 +3540,64 @@ TEST_F(ShapeInferenceTest, Pad) {
               HasSubstr("negative size for dimension 1"));
 }
 
+TEST_F(ShapeInferenceTest, PadAddsConstantOffsetToExpressions) {
+  const Shape input_shape = ShapeUtil::MakeShape(
+      F32, {4, 5}, std::vector<DExpr>{DExpr::Var(1), DExpr::Var(2)});
+  const Shape padding_value_shape = ShapeUtil::MakeShape(F32, {});
+  PaddingConfig padding_config;
+  auto* dimension0 = padding_config.add_dimensions();
+  dimension0->set_edge_padding_low(1);
+  dimension0->set_edge_padding_high(2);
+  dimension0->set_interior_padding(1);
+  auto* dimension1 = padding_config.add_dimensions();
+  dimension1->set_edge_padding_low(0);
+  dimension1->set_edge_padding_high(4);
+  dimension1->set_interior_padding(0);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferPadShape(input_shape, padding_value_shape,
+                                    padding_config));
+
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape, ShapeUtil::MakeShape(
+                          F32, {10, 9},
+                          std::vector<DExpr>{DExpr::Var(1) + 6,
+                                             DExpr::Var(2) + 4})));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(1) + 6));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Var(2) + 4));
+}
+
+TEST_F(ShapeInferenceTest, PadBuildsExpressionsForTwoSymbolicDimensions) {
+  const Shape input_shape = ShapeUtil::MakeShape(
+      F32, {7, 9}, std::vector<DExpr>{DExpr::Var(34), DExpr::Var(35)});
+  const Shape padding_value_shape = ShapeUtil::MakeShape(F32, {});
+  PaddingConfig padding_config;
+  auto* dimension0 = padding_config.add_dimensions();
+  dimension0->set_edge_padding_low(2);
+  dimension0->set_edge_padding_high(1);
+  dimension0->set_interior_padding(2);
+  auto* dimension1 = padding_config.add_dimensions();
+  dimension1->set_edge_padding_low(3);
+  dimension1->set_edge_padding_high(4);
+  dimension1->set_interior_padding(1);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferPadShape(input_shape, padding_value_shape,
+                                    padding_config));
+
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape, ShapeUtil::MakeShape(
+                          F32, {22, 24},
+                          std::vector<DExpr>{DExpr::Var(34) + 15,
+                                             DExpr::Var(35) + 15})));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             DExpr::Var(34) + 15));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1),
+                             DExpr::Var(35) + 15));
+}
+
 TEST_F(ShapeInferenceTest, Reverse) {
   const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25});
 
@@ -2997,6 +3607,20 @@ TEST_F(ShapeInferenceTest, Reverse) {
   ASSERT_TRUE(ShapeUtil::Equal(input_shape, *inferred_shape));
 }
 
+TEST_F(ShapeInferenceTest, ReversePreservesExpressions) {
+  const Shape input_shape = ShapeUtil::MakeShape(
+      F32, {10, 25, 7},
+      std::vector<DExpr>{DExpr::Var(26), DExpr::Const(25), DExpr::Var(27)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReverseShape(input_shape, {0, 2}));
+
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(26)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(25)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Var(27)));
+}
+
 TEST_F(ShapeInferenceTest, ReverseInvalidDimension) {
   const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25});
 
@@ -3071,6 +3695,21 @@ TEST_F(ShapeInferenceTest, Transpose) {
                                     *inferred_shape_and_status));
 }
 
+TEST_F(ShapeInferenceTest, TransposePermutesExpressions) {
+  const Shape a_shape = ShapeUtil::MakeShape(
+      F32, {2, 3, 4, 5},
+      std::vector<DExpr>{DExpr::Var(28), DExpr::Const(3), DExpr::Var(29),
+                         DExpr::Const(5)});
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferTransposeShape(
+                              a_shape, {1, 2, 3, 0}));
+
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Const(3)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Var(29)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(5)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(3), DExpr::Var(28)));
+}
+
 TEST_F(ShapeInferenceTest, Rank1Transpose) {
   const Shape a_shape = ShapeUtil::MakeShape(F32, {5});
   const absl::StatusOr<Shape> inferred_shape_and_status =
@@ -3260,8 +3899,8 @@ TEST_F(ShapeInferenceTest, ConditionalIndexed) {
 
 TEST_F(ShapeInferenceTest, ConditionalDynamic) {
   const Shape r0s32 = ShapeUtil::MakeShape(S32, {});
-  const Shape static_shape = ShapeUtil::MakeShape(S32, {4}, {false});
-  const Shape dynamic_shape = ShapeUtil::MakeShape(S32, {4}, {true});
+  const Shape static_shape = ShapeUtil::MakeShape(S32, {4}, {false}, {});
+  const Shape dynamic_shape = ShapeUtil::MakeShape(S32, {4}, {true}, {});
   const absl::StatusOr<Shape> inferred_shape0 =
       ShapeInference::InferConditionalShape(
           r0s32,
@@ -3342,6 +3981,25 @@ TEST_F(ShapeInferenceTest, GoodTopK) {
                                      ShapeUtil::MakeShape(S32, {3, 4, 2})})));
 }
 
+TEST_F(ShapeInferenceTest, TopKPreservesLeadingExpressions) {
+  const Shape input = ShapeUtil::MakeShape(
+      F32, {3, 4, 5},
+      std::vector<DExpr>{DExpr::Var(7), DExpr::Const(4), DExpr::Var(8)});
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferTopKShape(input, /*k=*/2));
+
+  ASSERT_TRUE(inferred_shape.IsTuple());
+  ASSERT_EQ(inferred_shape.tuple_shapes_size(), 2);
+  const Shape& values = inferred_shape.tuple_shapes(0);
+  const Shape& indices = inferred_shape.tuple_shapes(1);
+  EXPECT_TRUE(DynExpr::equal(values.expressions(0), DExpr::Var(7)));
+  EXPECT_TRUE(DynExpr::equal(values.expressions(1), DExpr::Const(4)));
+  EXPECT_TRUE(DynExpr::equal(values.expressions(2), DExpr::Const(2)));
+  EXPECT_TRUE(DynExpr::equal(indices.expressions(0), DExpr::Var(7)));
+  EXPECT_TRUE(DynExpr::equal(indices.expressions(1), DExpr::Const(4)));
+  EXPECT_TRUE(DynExpr::equal(indices.expressions(2), DExpr::Const(2)));
+}
+
 TEST_F(ShapeInferenceTest, FailTopKLargeK) {
   const Shape input = ShapeUtil::MakeShape(F32, {3, 4, 5});
   const absl::StatusOr<Shape> statusor =
@@ -3492,7 +4150,7 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherEntireDimension) {
   TF_ASSERT_OK_AND_ASSIGN(
       const Shape gather_shape,
       ShapeInference::InferGatherShape(
-          ShapeUtil::MakeShape(F32, {3, 2, 1}, {false, true, false}),
+          ShapeUtil::MakeShape(F32, {3, 2, 1}, {false, true, false}, {}),
           ShapeUtil::MakeShape(S64, {}),
           HloGatherInstruction::MakeGatherDimNumbers(
               /*offset_dims=*/{0, 1},
@@ -3501,7 +4159,8 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherEntireDimension) {
               /*index_vector_dim=*/0),
           /*slice_sizes=*/{1, 2, 1}));
   EXPECT_TRUE(ShapeUtil::Equal(
-      gather_shape, ShapeUtil::MakeShape(F32, {2, 1}, {true, false})))
+      gather_shape,
+      ShapeUtil::MakeShape(F32, {2, 1}, {true, false}, {})))
       << ShapeUtil::HumanString(gather_shape);
 }
 
@@ -3509,7 +4168,7 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherCollapsedDimension) {
   TF_ASSERT_OK_AND_ASSIGN(
       const Shape gather_shape,
       ShapeInference::InferGatherShape(
-          ShapeUtil::MakeShape(F32, {3, 2, 1}, {true, false, false}),
+          ShapeUtil::MakeShape(F32, {3, 2, 1}, {true, false, false}, {}),
           ShapeUtil::MakeShape(S64, {}),
           HloGatherInstruction::MakeGatherDimNumbers(
               /*offset_dims=*/{0, 1},
@@ -3518,7 +4177,8 @@ TEST_F(GatherShapeInferenceTest, DynamicGatherCollapsedDimension) {
               /*index_vector_dim=*/0),
           /*slice_sizes=*/{1, 2, 1}));
   EXPECT_TRUE(ShapeUtil::Equal(
-      gather_shape, ShapeUtil::MakeShape(F32, {2, 1}, {false, false})))
+      gather_shape,
+      ShapeUtil::MakeShape(F32, {2, 1}, {false, false}, {})))
       << ShapeUtil::HumanString(gather_shape);
 }
 
@@ -3527,7 +4187,7 @@ TEST_F(GatherShapeInferenceTest, DynamicIndices) {
       const Shape gather_shape,
       ShapeInference::InferGatherShape(
           ShapeUtil::MakeShape(F32, {3, 2, 2}),
-          ShapeUtil::MakeShape(S64, {3, 4, 2}, {false, true, false}),
+          ShapeUtil::MakeShape(S64, {3, 4, 2}, {false, true, false}, {}),
           HloGatherInstruction::MakeGatherDimNumbers(
               /*offset_dims=*/{2, 3},
               /*collapsed_slice_dims=*/{0},
@@ -3536,10 +4196,35 @@ TEST_F(GatherShapeInferenceTest, DynamicIndices) {
           /*slice_sizes=*/{1, 2, 2}));
   EXPECT_TRUE(ShapeUtil::Equal(
       gather_shape,
-      ShapeUtil::MakeShape(F32, {3, 4, 2, 2}, {false, true, false, false})))
+      ShapeUtil::MakeShape(F32, {3, 4, 2, 2}, {false, true, false, false},
+                           {})))
       << ShapeUtil::HumanString(gather_shape);
 }
 
+TEST_F(GatherShapeInferenceTest, GatherPreservesIndexAndSliceExpressions) {
+  const Shape input = ShapeUtil::MakeShape(
+      F32, {3, 2, 2},
+      std::vector<DExpr>{DExpr::Const(3), DExpr::Var(22), DExpr::Const(2)});
+  const Shape indices = ShapeUtil::MakeShape(
+      S64, {3, 4, 2}, std::vector<bool>{false, true, false},
+      std::vector<DExpr>{DExpr::Const(3), DExpr::Var(23), DExpr::Const(2)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape gather_shape,
+      ShapeInference::InferGatherShape(
+          input, indices,
+          HloGatherInstruction::MakeGatherDimNumbers(
+              /*offset_dims=*/{2, 3},
+              /*collapsed_slice_dims=*/{0},
+              /*start_index_map=*/{0, 1},
+              /*index_vector_dim=*/2),
+          /*slice_sizes=*/{1, 2, 2}));
+  EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(0), DExpr::Const(3)));
+  EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(1), DExpr::Var(23)));
+  EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(2), DExpr::Var(22)));
+  EXPECT_TRUE(DynExpr::equal(gather_shape.expressions(3), DExpr::Const(2)));
+}
+
 TEST_F(GatherShapeInferenceTest, NonDefaultGatherIndicesLeafDim_A) {
   TF_ASSERT_OK_AND_ASSIGN(
       const Shape gather_shape,
@@ -4708,6 +5393,20 @@ TEST_F(ShapeInferenceTest, UnboundedAllToAll) {
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, AllToAllPreservesExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {12, 5},
+      std::vector<DExpr>{DExpr::Var(14), DExpr::Const(5)});
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferAllToAllShape(/*shape=*/operand,
+                                         /*split_dimension=*/0,
+                                         /*concat_dimension=*/0,
+                                         /*split_count=*/3));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(14)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(5)));
+}
+
 TEST_F(ShapeInferenceTest, UnboundedAllToAllTupleUnsupported) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape expected,
@@ -4787,6 +5486,30 @@ TEST_F(ShapeInferenceTest, UnboundedBatchNormGrad) {
       << " expected: " << ShapeUtil::HumanString(expected_tuple_shape);
 }
 
+TEST_F(ShapeInferenceTest, BatchNormGradPreservesFeatureExpression) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {5, 7, 11}, std::vector<bool>{false, true, false},
+      std::vector<DExpr>{DExpr::Const(5), DExpr::Var(12), DExpr::Const(11)});
+  const Shape scale =
+      ShapeUtil::MakeShape(F32, {7}, std::vector<DExpr>{DExpr::Var(12)});
+  const Shape mean =
+      ShapeUtil::MakeShape(F32, {7}, std::vector<DExpr>{DExpr::Var(12)});
+  const Shape variance =
+      ShapeUtil::MakeShape(F32, {7}, std::vector<DExpr>{DExpr::Var(12)});
+  const Shape output_grad = operand;
+
+  TF_ASSERT_OK_AND_ASSIGN(const Shape inferred_shape,
+                          ShapeInference::InferBatchNormGradShape(
+                              operand, scale, mean, variance, output_grad, 1));
+  ASSERT_TRUE(inferred_shape.IsTuple());
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(1), DExpr::Var(12)));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(0), DExpr::Var(12)));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(2).expressions(0), DExpr::Var(12)));
+}
+
 TEST_F(ShapeInferenceTest, UnboundedBatchNormInference) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, ?, 7]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape scale, ParseShape("f32[5]"));
@@ -4819,6 +5542,27 @@ TEST_F(ShapeInferenceTest, UnboundedBatchNormTraining) {
       << " expected: " << ShapeUtil::HumanString(expected_tuple_shape);
 }
 
+TEST_F(ShapeInferenceTest, BatchNormTrainingPreservesFeatureExpression) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {5, 7, 11}, std::vector<bool>{false, true, false},
+      std::vector<DExpr>{DExpr::Const(5), DExpr::Var(13), DExpr::Const(11)});
+  const Shape scale =
+      ShapeUtil::MakeShape(F32, {7}, std::vector<DExpr>{DExpr::Var(13)});
+  const Shape offset =
+      ShapeUtil::MakeShape(F32, {7}, std::vector<DExpr>{DExpr::Var(13)});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferBatchNormTrainingShape(operand, scale, offset, 1));
+  ASSERT_TRUE(inferred_shape.IsTuple());
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(1), DExpr::Var(13)));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(0), DExpr::Var(13)));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(2).expressions(0), DExpr::Var(13)));
+}
+
 TEST_F(ShapeInferenceTest, UnboundedBroadcastUnsupportedOperand) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[<=2, ?]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[1, <=2, ?]"));
@@ -5208,6 +5952,29 @@ TEST_F(ShapeInferenceTest, UnboundedDotGeneral) {
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, DotGeneralPreservesBatchExpression) {
+  const Shape lhs = ShapeUtil::MakeShape(
+      F32, {2, 3, 5}, std::vector<bool>{true, false, false},
+      {DExpr::Var(1), DExpr::Const(3), DExpr::Const(5)});
+  const Shape rhs = ShapeUtil::MakeShape(
+      F32, {2, 5, 7}, std::vector<bool>{true, false, false},
+      {DExpr::Var(1), DExpr::Const(5), DExpr::Const(7)});
+
+  DotDimensionNumbers dnums;
+  dnums.add_lhs_batch_dimensions(0);
+  dnums.add_rhs_batch_dimensions(0);
+  dnums.add_lhs_contracting_dimensions(2);
+  dnums.add_rhs_contracting_dimensions(1);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferDotOpShape(lhs, rhs, dnums,
+                                      /*preferred_element_type=*/std::nullopt));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(1)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(3)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(2), DExpr::Const(7)));
+}
+
 TEST_F(ShapeInferenceTest, UnboundedDynamicSlice) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape start_index, ParseShape("s32[]"));
@@ -5216,12 +5983,30 @@ TEST_F(ShapeInferenceTest, UnboundedDynamicSlice) {
       const Shape inferred_shape,
       ShapeInference::InferDynamicSliceShape(
           operand, /*start_index_shapes=*/{start_index, start_index},
-          /*slice_sizes=*/{2, 2}, /*allow_scalar_indices=*/true));
+          /*slice_sizes=*/{2, 2}, /*slice_exprs=*/{},
+          /*allow_scalar_indices=*/true));
   EXPECT_TRUE(ShapeUtil::Equal(inferred_shape, expected))
       << "inferred: " << ShapeUtil::HumanString(inferred_shape)
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, DynamicSliceUsesProvidedExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {9, 10},
+      std::vector<DExpr>{DExpr::Var(15), DExpr::Const(10)});
+  const Shape start_index = ShapeUtil::MakeShape(S32, {});
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferDynamicSliceShape(
+          operand, /*start_index_shapes=*/{start_index, start_index},
+          /*slice_sizes=*/{4, 10},
+          /*slice_exprs=*/{DExpr::Var(15) / 2, DExpr::Const(10)},
+          /*allow_scalar_indices=*/true));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(15) / 2));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(10)));
+}
+
 TEST_F(ShapeInferenceTest, UnboundedDynamicUpdateSlice) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape update, ParseShape("f32[?, 5]"));
@@ -5237,6 +6022,42 @@ TEST_F(ShapeInferenceTest, UnboundedDynamicUpdateSlice) {
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, DynamicUpdateSlicePreservesOperandExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {12, 10},
+      std::vector<DExpr>{DExpr::Var(30), DExpr::Const(10)});
+  const Shape update = ShapeUtil::MakeShape(
+      F32, {4, 10},
+      std::vector<DExpr>{DExpr::Const(4), DExpr::Const(10)});
+  const Shape start_index = ShapeUtil::MakeShape(S32, {});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferDynamicUpdateSliceShape(
+          operand, update, /*start_index_shapes=*/{start_index, start_index},
+          /*allow_scalar_indices=*/true));
+
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(30)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Const(10)));
+}
+
+TEST_F(ShapeInferenceTest, DynamicReshapePreservesExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {5, 4, 8},
+      std::vector<bool>{false, false, false},
+      std::vector<DExpr>{DExpr::Var(21), DExpr::Const(4), DExpr::Const(8)});
+  const Shape dim_size = ShapeUtil::MakeShape(S32, {});
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferDynamicReshapeShape(
+          operand, /*dim_size_shapes=*/{&dim_size},
+          /*new_size_bounds=*/{160},
+          /*dims_are_dynamic=*/{false},
+          /*expressions=*/{DExpr::Var(21) * DExpr::Const(32)}));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0),
+                             DExpr::Var(21) * DExpr::Const(32)));
+}
+
 TEST_F(ShapeInferenceTest, UnboundedFftWithFFT) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("c64[2, <=5, ?]"));
   const std::vector<int64_t> fft_length = {5, 10};
@@ -5472,6 +6293,56 @@ TEST_F(ShapeInferenceTest, UnboundedReduce) {
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, ReducePreservesRemainingExpressions) {
+  const Shape input = ShapeUtil::MakeShape(
+      F32, {5, 7, 11},
+      std::vector<DExpr>{DExpr::Var(9), DExpr::Const(7), DExpr::Var(10)});
+  ProgramShape to_apply =
+      ShapeUtil::MakeProgramShape({f32_, f32_}, f32_);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReduceShape({&input, &f32_}, {1}, to_apply));
+
+  EXPECT_TRUE(ShapeUtil::Equal(
+      inferred_shape, ShapeUtil::MakeShape(
+                          F32, {5, 11},
+                          std::vector<DExpr>{DExpr::Var(9), DExpr::Var(10)})));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(0), DExpr::Var(9)));
+  EXPECT_TRUE(DynExpr::equal(inferred_shape.expressions(1), DExpr::Var(10)));
+}
+
+TEST_F(ShapeInferenceTest, ReduceTupleOutputsPreserveRemainingExpressions) {
+  const Shape input0 = ShapeUtil::MakeShape(
+      F32, {5, 7, 11},
+      std::vector<DExpr>{DExpr::Var(31), DExpr::Const(7), DExpr::Var(32)});
+  const Shape input1 = ShapeUtil::MakeShape(
+      S32, {5, 7, 11},
+      std::vector<DExpr>{DExpr::Var(31), DExpr::Const(7), DExpr::Var(32)});
+  ProgramShape to_apply = ShapeUtil::MakeProgramShape(
+      {f32_, s32_, f32_, s32_}, ShapeUtil::MakeTupleShape({f32_, s32_}));
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape inferred_shape,
+      ShapeInference::InferReduceShape(
+          {&input0, &input1, &f32_, &s32_}, {1}, to_apply));
+
+  ASSERT_TRUE(inferred_shape.IsTuple());
+  ASSERT_EQ(inferred_shape.tuple_shapes_size(), 2);
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(0),
+                     DExpr::Var(31)));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(0).expressions(1),
+                     DExpr::Var(32)));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(0),
+                     DExpr::Var(31)));
+  EXPECT_TRUE(
+      DynExpr::equal(inferred_shape.tuple_shapes(1).expressions(1),
+                     DExpr::Var(32)));
+}
+
 TEST_F(ShapeInferenceTest, UnboundedReduceInvalidReduceDimension) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape input0, ParseShape("f32[7, 5]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape input1, ParseShape("f32[?, 5]"));
@@ -5562,7 +6433,8 @@ TEST_F(ShapeInferenceTest, UnboundedReshape) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[2,3]"));
   TF_ASSERT_OK_AND_ASSIGN(
       const Shape inferred,
-      ShapeInference::InferReshapeShape(operand, /*dimensions=*/{2, 3}, -1));
+      ShapeInference::InferReshapeShape(operand, /*dimensions=*/{2, 3}, -1,
+                                        /*expressions=*/{}));
   ASSERT_TRUE(ShapeUtil::Equal(inferred, expected))
       << "inferred: " << ShapeUtil::HumanString(inferred)
       << " expected: " << ShapeUtil::HumanString(expected);
@@ -5573,7 +6445,8 @@ TEST_F(ShapeInferenceTest, UnboundedReshapeUnsupportedOutputShape) {
   const absl::StatusOr<Shape> inferred_shape =
       ShapeInference::InferReshapeShape(
           operand,
-          /*dimensions=*/{Shape::kUnboundedSize, Shape::kUnboundedSize}, -1);
+          /*dimensions=*/{Shape::kUnboundedSize, Shape::kUnboundedSize}, -1,
+          /*expressions=*/{});
   EXPECT_THAT(
       inferred_shape.status().message(),
       HasSubstr("Reshaping with unbounded result shape is not supported."));
@@ -5583,7 +6456,8 @@ TEST_F(ShapeInferenceTest, UnboundedReshapeUnsupportedMixOfDynamism) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, <=3]"));
   TF_ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[<=3]"));
   const absl::StatusOr<Shape> inferred_shape =
-      ShapeInference::InferReshapeShape(operand, /*dimensions=*/{3}, -1);
+      ShapeInference::InferReshapeShape(operand, /*dimensions=*/{3}, -1,
+                                        /*expressions=*/{});
   ASSERT_THAT(inferred_shape.status().message(),
               HasSubstr("Reshape operand with bounded and unbounded dynamism "
                         "not supported."));
@@ -5693,6 +6567,47 @@ TEST_F(ShapeInferenceTest, UnboundedSelectAndScatter) {
       << " expected: " << ShapeUtil::HumanString(expected);
 }
 
+TEST_F(ShapeInferenceTest, SelectAndScatterPreservesOperandExpressions) {
+  const Shape operand = ShapeUtil::MakeShape(
+      F32, {11, 10}, std::vector<DExpr>{DExpr::Var(36), DExpr::Const(10)});
+  const Shape source = ShapeUtil::MakeShape(
+      F32, {5, 10}, std::vector<DExpr>{(((DExpr::Var(36) - 1) / 2) + 1).simplify(),
+                                       DExpr::Const(10)});
+  const Shape init_value = ShapeUtil::MakeShape(F32, {});
+
+  Window window;
+  WindowDimension dim0;
+  dim0.set_base_dilation(1);
+  dim0.set_size(3);
+  dim0.set_stride(2);
+  dim0.set_padding_low(0);
+  dim0.set_padding_high(1);
+  dim0.set_window_dilation(1);
+
+  WindowDimension dim1;
+  dim1.set_base_dilation(1);
+  dim1.set_size(1);
+  dim1.set_stride(1);
+  dim1.set_padding_low(0);
+  dim1.set_padding_high(0);
+  dim1.set_window_dilation(1);
+
+  *window.add_dimensions() = dim0;
+  *window.add_dimensions() = dim1;
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      const Shape result,
+      ShapeInference::InferSelectAndScatterShape(
+          operand,
+          /*select_shape=*/ShapeUtil::MakeProgramShape({f32_, f32_}, pred_),
+          window, source, init_value,
+          /*scatter_shape=*/
+          ShapeUtil::MakeProgramShape({f32_, f32_}, f32_)));
+
+  EXPECT_TRUE(DynExpr::equal(result.expressions(0), DExpr::Var(36)));
+  EXPECT_TRUE(DynExpr::equal(result.expressions(1), DExpr::Const(10)));
+}
+
 TEST_P(UnboundedBinaryOpShapeInferenceTest, UnboundedShiftLeft) {
   TF_ASSERT_OK_AND_ASSIGN(const Shape lhs, ParseShape(GetParam().lhs));
   TF_ASSERT_OK_AND_ASSIGN(const Shape rhs, ParseShape(GetParam().rhs));
diff --git a/third_party/xla/xla/service/triangular_solve_expander.cc b/third_party/xla/xla/service/triangular_solve_expander.cc
index 67dc0235810de8..3e51ffd4de8439 100644
--- a/third_party/xla/xla/service/triangular_solve_expander.cc
+++ b/third_party/xla/xla/service/triangular_solve_expander.cc
@@ -120,7 +120,12 @@ XlaOp DiagonalBlocks(XlaOp a, int64_t block_size) {
       auto last_blocks_dims = std::vector<int64_t>(ndims);
       std::copy(shape_dims.begin(), shape_dims.end(), last_blocks_dims.begin());
       last_blocks_dims.insert(last_blocks_dims.end() - 2, 1);
-      last_blocks = Reshape(last_blocks, last_blocks_dims);
+      auto shape_exprs = blocks_shape.expressions();
+      auto last_blocks_exprs = std::vector<DExpr>(ndims);
+      std::copy(shape_exprs.begin(), shape_exprs.end(),
+                last_blocks_exprs.begin());
+      last_blocks_exprs.insert(last_blocks_exprs.end() - 2, DExpr::Const(1));
+      last_blocks = Reshape(last_blocks, last_blocks_dims, last_blocks_exprs);
 
       // Concatenate with the other blocks if necessary
       if (n > block_size) {
@@ -366,7 +371,7 @@ XlaOp TriangularSolveExpander::InvertDiagonalBlocks(
                           /*broadcast_dimensions=*/{0, 1});
 
     // Reshape back to original batch major dimensions
-    return Reshape(inv_diag_blocks, shape.dimensions());
+    return Reshape(inv_diag_blocks, shape.dimensions(), shape.expressions());
   });
 }
 
diff --git a/third_party/xla/xla/shape.cc b/third_party/xla/xla/shape.cc
index 1cee38146fb07d..e12d07a25ecfe5 100644
--- a/third_party/xla/xla/shape.cc
+++ b/third_party/xla/xla/shape.cc
@@ -17,6 +17,9 @@ limitations under the License.
 
 #include <algorithm>
 #include <cstdint>
+#include <cstdlib>
+#include <map>
+#include <numeric>
 #include <optional>
 #include <ostream>
 #include <string>
@@ -40,7 +43,6 @@ limitations under the License.
 #include "xla/xla_data.pb.h"
 
 namespace xla {
-
 // Defined in .cc file to avoid inlining these large routines
 Shape::Shape() = default;
 Shape::~Shape() = default;
@@ -97,13 +99,18 @@ Shape::Shape(const ShapeProto& shape_proto) {
 }
 
 absl::StatusOr<Shape> Shape::FromProto(const ShapeProto& shape_proto) {
+
+  // LOG(INFO) << "FROM PROTO:\n" << shape_proto.DebugString() << std::endl;
+
   Shape shape;
   shape.set_element_type(shape_proto.element_type());
   if (auto* const state = shape.if_array_state()) {
     const int num_dims = shape_proto.dimensions_size();
     const int num_is_dynamic_dims = shape_proto.is_dynamic_dimension_size();
+    const int num_expressions = shape_proto.expressions_size();
     state->dimensions.reserve(num_dims);
     state->dynamic_dimensions.reserve(num_dims);
+    state->expressions.reserve(num_dims);
     if (num_is_dynamic_dims != 0) {
       TF_RET_CHECK(num_dims == num_is_dynamic_dims)
           << "Malformed shape proto: number of is_dynamic_dimension "
@@ -111,6 +118,13 @@ absl::StatusOr<Shape> Shape::FromProto(const ShapeProto& shape_proto) {
           << num_is_dynamic_dims << ") does not match number of dimension "
           << "fields (" << num_dims << ").";
     }
+    if (num_expressions != 0) {
+      TF_RET_CHECK(num_dims == num_expressions)
+          << "Malformed shape proto: number of expressions "
+             "fields ("
+          << num_expressions << ") does not match number of dimension "
+          << "fields (" << num_dims << ").";
+    }
     for (int i = 0; i < num_dims; ++i) {
       const bool is_dynamic =
           (i < num_is_dynamic_dims) && shape_proto.is_dynamic_dimension(i);
@@ -118,7 +132,11 @@ absl::StatusOr<Shape> Shape::FromProto(const ShapeProto& shape_proto) {
       // UnsafeAddDimension. We expect that the caller will eventually call a
       // validation routine that will detect the error in case the dimension
       // value is invalid.
-      shape.UnsafeAddDimension(shape_proto.dimensions(i), is_dynamic);
+      DExpr expression =
+          (i < num_expressions) ? DExprFromProto(shape_proto.expressions(i))
+                                : DExpr::Const(shape_proto.dimensions(i));
+      shape.UnsafeAddDimension(shape_proto.dimensions(i), is_dynamic,
+                               expression);
     }
   } else if (auto* const state = shape.if_tuple_state()) {
     state->tuple_shapes.reserve(shape_proto.tuple_shapes_size());
@@ -136,6 +154,7 @@ absl::StatusOr<Shape> Shape::FromProto(const ShapeProto& shape_proto) {
     TF_ASSIGN_OR_RETURN(*shape.mutable_layout(),
                         Layout::FromProto(shape_proto.layout()));
   }
+  // LOG(INFO) << "FROM PROTO " << shape << "\n";
   return shape;
 }
 
@@ -143,6 +162,8 @@ ShapeProto Shape::ToProto() const {
   ShapeProto proto;
   proto.set_element_type(element_type_);
 
+  // LOG(INFO) << "TO PROTO " << ToString() << "\n";
+
   if (const auto* const state = if_array_state()) {
     proto.mutable_dimensions()->Reserve(state->dimensions.size());
     for (const int64_t dimension : state->dimensions) {
@@ -151,6 +172,11 @@ ShapeProto Shape::ToProto() const {
     for (const bool dynamic : state->dynamic_dimensions) {
       proto.add_is_dynamic_dimension(dynamic);
     }
+    for (const DExpr& e : state->expressions) {
+      ExpressionProto* eproto = proto.add_expressions();
+      CHECK(e.get() != nullptr) << "Missing expression in expression list.";
+      e->to_proto(eproto);
+    }
     if (state->layout.has_value()) {
       *proto.mutable_layout() = state->layout->ToProto();
     }
@@ -163,6 +189,7 @@ ShapeProto Shape::ToProto() const {
     proto.mutable_tuple_shapes()->Reserve(1);
     *proto.add_tuple_shapes() = state->buffer_shape[0].ToProto();
   }
+  // LOG(INFO) << "DEBUG VIEW:\n" << proto.DebugString() << std::endl;
   return proto;
 }
 
@@ -245,14 +272,16 @@ bool Shape::AreAllLeavesIntegers() const {
   return primitive_util::IsIntegralType(element_type());
 }
 
-void Shape::add_dimensions(int64_t value, bool is_dynamic) {
+void Shape::add_dimensions(int64_t value, bool is_dynamic, DExpr expr) {
   if (value < 0) {
     CHECK(is_dynamic) << "static dimension must have size >= 0 instead of "
                       << value << ".";
     CHECK_EQ(value, kUnboundedSize)
         << "dynamic dimension must have size == kUnboundedSize or >= 0.";
   }
-  UnsafeAddDimension(value, is_dynamic);
+  UnsafeAddDimension(
+      value, is_dynamic,
+      expr ? std::move(expr) : DExpr::Const(value));
 }
 
 void Shape::set_dynamic_dimension(int dimension, bool is_dynamic) {
@@ -262,6 +291,26 @@ void Shape::set_dynamic_dimension(int dimension, bool is_dynamic) {
   state.dynamic_dimensions[dimension] = is_dynamic;
 }
 
+void Shape::set_expression(int dimension, DExpr e) {
+  auto& state = array_state();
+  state.expressions[dimension] =
+      e ? std::move(e) : DExpr::Const(state.dimensions[dimension]);
+}
+
+void Shape::set_expressions(std::vector<DExpr> exps) {
+  auto& state = array_state();
+  CHECK_LE(exps.size(), state.dimensions.size());
+  state.expressions.resize(state.dimensions.size());
+  for (size_t i = 0; i < state.dimensions.size(); ++i) {
+    state.expressions[i] = i < exps.size()
+                               ? std::move(exps[i])
+                               : DExpr::Const(state.dimensions[i]);
+    if (!state.expressions[i]) {
+      state.expressions[i] = DExpr::Const(state.dimensions[i]);
+    }
+  }
+}
+
 void Shape::set_dimensions(int index, int64_t size,
                            std::optional<bool> is_dynamic) {
   auto& state = array_state();
@@ -270,6 +319,7 @@ void Shape::set_dimensions(int index, int64_t size,
   CheckDimensionSize(index, size, dynamic);
   state.dimensions[index] = size;
   state.dynamic_dimensions[index] = dynamic;
+  state.expressions[index] = DExpr::Const(size);
 }
 
 void Shape::set_dimensions_minor(int index, int64_t size,
@@ -291,12 +341,15 @@ void Shape::CheckDimensionSize(int dim_index, int64_t size, bool is_dynamic) {
   }
 }
 
-void Shape::UnsafeAddDimension(int64_t value, bool is_dynamic) {
+void Shape::UnsafeAddDimension(int64_t value, bool is_dynamic, DExpr exp) {
   auto& state = array_state();
   CHECK_EQ(state.dimensions.size(), state.dynamic_dimensions.size())
       << "where the shape is " << ToString();
+  CHECK_EQ(state.dimensions.size(), state.expressions.size())
+      << "where the shape is " << ToString();
   state.dimensions.push_back(value);
   state.dynamic_dimensions.push_back(is_dynamic);
+  state.expressions.push_back(exp ? std::move(exp) : DExpr::Const(value));
 }
 
 bool Shape::is_static() const {
@@ -345,6 +398,8 @@ void Shape::DeleteDimension(int64_t dim_to_delete) {
   state.dimensions.erase(state.dimensions.begin() + dim_to_delete);
   state.dynamic_dimensions.erase(state.dynamic_dimensions.begin() +
                                  dim_to_delete);
+  state.expressions.erase(state.expressions.begin() +
+                                 dim_to_delete);
   if (LayoutUtil::HasLayout(*this)) {
     state.layout->DeleteDimension(dim_to_delete);  // NOLINT: optional-access
   }
@@ -358,6 +413,8 @@ void Shape::DeleteDimensions(absl::Span<const int64_t> dims_to_delete) {
   state.dimensions = RemoveElements(sorted_dims_to_delete, state.dimensions);
   state.dynamic_dimensions =
       RemoveElements(sorted_dims_to_delete, state.dynamic_dimensions);
+  state.expressions =
+      RemoveElements(sorted_dims_to_delete, state.expressions);
   if (LayoutUtil::HasLayout(*this)) {
     for (auto it = sorted_dims_to_delete.rbegin();
          it != sorted_dims_to_delete.rend(); ++it) {
@@ -370,6 +427,7 @@ void Shape::CheckStateIsEmpty() const {
   if (const auto* const state = if_array_state()) {
     CHECK(state->dimensions.empty()) << ToString();
     CHECK(state->dynamic_dimensions.empty()) << ToString();
+    CHECK(state->expressions.empty()) << ToString();
     CHECK(!state->layout.has_value()) << ToString();
   } else if (const auto* const state = if_tuple_state()) {
     CHECK(state->tuple_shapes.empty()) << ToString();
@@ -509,6 +567,11 @@ bool Shape::Equal::operator()(const Shape& lhs, const Shape& rhs) {
            rhs.is_unbounded_dynamic_dimension(i))) {
         continue;
       }
+      if (i == 0 && ignore_batch_ &&
+          (lhs.outer_multiplier() > 0 || rhs.outer_multiplier() > 0)) {
+        VLOG(3) << "CompareShapes: batch dimension found. Forcely compatible";
+        continue;
+      }
       if (lhs.dimensions(i) != rhs.dimensions(i)) {
         VLOG(3) << "CompareShapes: lhs dimensions != rhs dimensions";
         return false;
diff --git a/third_party/xla/xla/shape.h b/third_party/xla/xla/shape.h
index 8453dc17717e11..13079dab62532d 100644
--- a/third_party/xla/xla/shape.h
+++ b/third_party/xla/xla/shape.h
@@ -38,6 +38,7 @@ limitations under the License.
 #include "xla/tsl/platform/logging.h"  // IWYU pragma: keep
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
+#include "xla/shape_expr.h"
 
 namespace xla {
 
@@ -218,6 +219,29 @@ class Shape {
     return array_state().dynamic_dimensions[dimension];
   }
 
+  bool has_dynamic_expr() const {
+    if (auto* const state = if_array_state()) {
+      return absl::c_any_of(state->expressions,
+                            [](const DExpr& e) {
+                              return e && e->is_dynamic();
+                            });
+    }
+    if (auto* const state = if_tuple_state()) {
+      return absl::c_any_of(state->tuple_shapes, [](Shape subshape) {
+        return subshape.has_dynamic_expr();
+      });
+    }
+    return false;
+  }
+
+  const DExpr& expressions(int dimension) const {
+    if (dimension < 0) return MissingExpression();
+    const auto& exprs = array_state().expressions;
+    const size_t dim = static_cast<size_t>(dimension);
+    if (dim >= exprs.size()) return MissingExpression();
+    return exprs[dim] ? exprs[dim] : MissingExpression();
+  }
+
   // Returns true if the given dimension is statically-sized.
   // Precondition: this is an array shape and `dimension` is a valid dimension
   // index.
@@ -232,12 +256,18 @@ class Shape {
   //   - The dimension's size is valid for the given dynamic-ness.
   void set_dynamic_dimension(int dimension, bool is_dynamic);
 
+  void set_expression(int dimension, DExpr e);
+
+  void set_expressions(std::vector<DExpr> exprs);
+
   // Returns a span to indicate whether each dimension is dynamic.
   // Precondition: this is an array shape.
   absl::Span<const bool> dynamic_dimensions() const {
     return array_state().dynamic_dimensions;
   }
 
+  absl::Span<const DExpr> expressions() const { return array_state().expressions; }
+
   // Removes the given dimension from the shape. Layout, if it exists, is
   // adjusted to match the modified shape.
   // Precondition: this is an array shape, and the input dimension indices are
@@ -313,7 +343,8 @@ class Shape {
   //   - This is an array shape.
   //   - Either `value` is >= 0, or `is_dynamic` is true and `value` is
   //     kUnboundedSize.
-  void add_dimensions(int64_t value, bool is_dynamic = false);
+  void add_dimensions(int64_t value, bool is_dynamic = false,
+                      DExpr expr = DExpr());
 
   // Clears all dimensions (i.e. makes this shape a scalar).
   // Precondition: this is an array shape.
@@ -321,6 +352,7 @@ class Shape {
     auto& state = array_state();
     state.dimensions.clear();
     state.dynamic_dimensions.clear();
+    state.expressions.clear();
   }
 
   // Returns a span to indicate the size of each dimension.
@@ -434,6 +466,10 @@ class Shape {
 
     bool operator()(const Shape& lhs, const Shape& rhs);
 
+    Equal& IgnoreBatch(bool ignore_batch = true) {
+      ignore_batch_ = ignore_batch;
+      return *this;
+    }
     Equal& IgnoreLayout(bool ignore_layout = true) {
       ignore_layout_ = ignore_layout;
       return *this;
@@ -488,6 +524,7 @@ class Shape {
     }
 
    private:
+    bool ignore_batch_ = false;
     bool ignore_layout_ = false;
     bool ignore_tiles_in_layout_ = false;
     bool ignore_element_size_in_layout_ = false;
@@ -515,7 +552,7 @@ class Shape {
     }
     if (const auto* const state = s.if_array_state()) {
       h = H::combine(std::move(h), s.element_type_, state->dimensions,
-                     state->dynamic_dimensions);
+                     state->dynamic_dimensions, state->expressions);
       if (kIsLayoutSensitive) {
         h = H::combine(std::move(h), state->layout);
       }
@@ -532,7 +569,11 @@ class Shape {
     return Shape::Hash(std::move(h), s);
   }
 
+  int64_t outer_multiplier() const { return outer_multiplier_; }
+  void set_outer_multiplier(int64_t m) { outer_multiplier_ = m; }
  private:
+  int64_t outer_multiplier_ = -1;
+
   friend absl::Status ValidateNonLayoutProperties(const Shape& shape);
 
   // Define one state struct for each shape category. Depending on the element
@@ -560,8 +601,27 @@ class Shape {
     // respective dimension is dynamically sized.
     absl::InlinedVector<bool, InlineRank()> dynamic_dimensions;
 
+    absl::InlinedVector<DExpr, InlineRank()> expressions;
+
     // The layout of the shape.
     std::optional<Layout> layout;
+
+    ArrayState() = default;
+    ArrayState(const ArrayState& other)
+        : dimensions(other.dimensions),
+          dynamic_dimensions(other.dynamic_dimensions),
+          expressions(other.expressions),
+          layout(other.layout) {}
+    ArrayState& operator=(const ArrayState& other) {
+      if (this == &other) return *this;
+      dimensions = other.dimensions;
+      dynamic_dimensions = other.dynamic_dimensions;
+      expressions = other.expressions;
+      layout = other.layout;
+      return *this;
+    }
+    ArrayState(ArrayState&&) noexcept = default;
+    ArrayState& operator=(ArrayState&&) noexcept = default;
   };
   struct TupleState {
     // The tuple element subshapes.
@@ -578,12 +638,13 @@ class Shape {
 
   // CHECKs that the dimension size is valid.
   void CheckDimensionSize(int dim_index, int64_t size, bool is_dynamic);
+  static const DExpr& MissingExpression();
 
   // Like add_dimensions(), but does not CHECK that the arguments are valid.
   // Instead, we rely on validation down the road to catch invalid shapes.
   // This is useful for code that should not crash, such as constructing a
   // Shape from an unvalidated proto.
-  void UnsafeAddDimension(int64_t value, bool is_dynamic);
+  void UnsafeAddDimension(int64_t value, bool is_dynamic, DExpr exp);
 
   // Convenience accessors for the state_ variant. Each if_*_state() accessor
   // returns a pointer to the corresponding state struct, or nullptr if the
diff --git a/third_party/xla/xla/shape_expr.cc b/third_party/xla/xla/shape_expr.cc
new file mode 100644
index 00000000000000..d74a42175b3380
--- /dev/null
+++ b/third_party/xla/xla/shape_expr.cc
@@ -0,0 +1,773 @@
+/* Copyright 2018 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/shape.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+#include <map>
+#include <memory>
+#include <ostream>
+#include <numeric>
+#include <optional>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "xla/printer.h"
+
+namespace xla {
+
+namespace {
+
+Constant* AsConstant(DynExpr* expr) {
+  return expr != nullptr && expr->kind() == DExpr::Kind::kConstant
+             ? static_cast<Constant*>(expr)
+             : nullptr;
+}
+
+std::vector<DynExpr*> ExpressionChildren(DynExpr* expr) {
+  CHECK(expr != nullptr);
+  switch (expr->kind()) {
+    case DExpr::Kind::kUnknown:
+    case DExpr::Kind::kConstant:
+    case DExpr::Kind::kVariable:
+      return {};
+    case DExpr::Kind::kAdd: {
+      auto* add = static_cast<Add*>(expr);
+      return {add->get_lhs(), add->get_rhs()};
+    }
+    case DExpr::Kind::kSub: {
+      auto* sub = static_cast<Sub*>(expr);
+      return {sub->get_lhs(), sub->get_rhs()};
+    }
+    case DExpr::Kind::kMul: {
+      auto* mul = static_cast<Mul*>(expr);
+      return {mul->get_lhs(), mul->get_rhs()};
+    }
+    case DExpr::Kind::kDiv: {
+      auto* div = static_cast<Div*>(expr);
+      return {div->get_lhs(), div->get_rhs()};
+    }
+    case DExpr::Kind::kMax: {
+      auto* max = static_cast<MaxExpr*>(expr);
+      return {max->get_lhs(), max->get_rhs()};
+    }
+    case DExpr::Kind::kGt: {
+      auto* gt = static_cast<GtExpr*>(expr);
+      return {gt->get_lhs(), gt->get_rhs()};
+    }
+    case DExpr::Kind::kSelect: {
+      auto* select = static_cast<SelectExpr*>(expr);
+      return {select->get_pred(), select->get_on_true(),
+              select->get_on_false()};
+    }
+  }
+  return {};
+}
+
+DynExpr* FindSmallestCoveringSubexpression(DynExpr* expr) {
+  CHECK(expr != nullptr);
+  if (expr->kind() == DExpr::Kind::kVariable) {
+    return expr;
+  }
+
+  DynExpr* common_core = nullptr;
+  for (DynExpr* child : ExpressionChildren(expr)) {
+    DynExpr* child_core = FindSmallestCoveringSubexpression(child);
+    if (child_core == nullptr) {
+      continue;
+    }
+    if (common_core == nullptr) {
+      common_core = child_core;
+    } else if (!DynExpr::equal(common_core, child_core)) {
+      return expr;
+    }
+  }
+  return common_core;
+}
+
+std::unique_ptr<DynExpr> ReplaceSubexpression(DynExpr* expr, DynExpr* target,
+                                              DynExpr* replacement) {
+  CHECK(expr != nullptr);
+  CHECK(target != nullptr);
+  CHECK(replacement != nullptr);
+  if (DynExpr::equal(expr, target)) {
+    return replacement->clone();
+  }
+
+  switch (expr->kind()) {
+    case DExpr::Kind::kUnknown:
+    case DExpr::Kind::kConstant:
+    case DExpr::Kind::kVariable:
+      return expr->clone();
+    case DExpr::Kind::kAdd: {
+      auto* add = static_cast<Add*>(expr);
+      return std::make_unique<Add>(
+          ReplaceSubexpression(add->get_lhs(), target, replacement).release(),
+          ReplaceSubexpression(add->get_rhs(), target, replacement).release());
+    }
+    case DExpr::Kind::kSub: {
+      auto* sub = static_cast<Sub*>(expr);
+      return std::make_unique<Sub>(
+          ReplaceSubexpression(sub->get_lhs(), target, replacement).release(),
+          ReplaceSubexpression(sub->get_rhs(), target, replacement).release());
+    }
+    case DExpr::Kind::kMul: {
+      auto* mul = static_cast<Mul*>(expr);
+      return std::make_unique<Mul>(
+          ReplaceSubexpression(mul->get_lhs(), target, replacement).release(),
+          ReplaceSubexpression(mul->get_rhs(), target, replacement).release());
+    }
+    case DExpr::Kind::kDiv: {
+      auto* div = static_cast<Div*>(expr);
+      return std::make_unique<Div>(
+          ReplaceSubexpression(div->get_lhs(), target, replacement).release(),
+          ReplaceSubexpression(div->get_rhs(), target, replacement).release());
+    }
+    case DExpr::Kind::kMax: {
+      auto* max = static_cast<MaxExpr*>(expr);
+      return std::make_unique<MaxExpr>(
+          ReplaceSubexpression(max->get_lhs(), target, replacement).release(),
+          ReplaceSubexpression(max->get_rhs(), target, replacement).release());
+    }
+    case DExpr::Kind::kGt: {
+      auto* gt = static_cast<GtExpr*>(expr);
+      return std::make_unique<GtExpr>(
+          ReplaceSubexpression(gt->get_lhs(), target, replacement).release(),
+          ReplaceSubexpression(gt->get_rhs(), target, replacement).release());
+    }
+    case DExpr::Kind::kSelect: {
+      auto* select = static_cast<SelectExpr*>(expr);
+      return std::make_unique<SelectExpr>(
+          ReplaceSubexpression(select->get_pred(), target, replacement)
+              .release(),
+          ReplaceSubexpression(select->get_on_true(), target, replacement)
+              .release(),
+          ReplaceSubexpression(select->get_on_false(), target, replacement)
+              .release());
+    }
+  }
+  return expr->clone();
+}
+
+void NormalizeFraction(int64_t* numerator, int64_t* denominator) {
+  CHECK(denominator != nullptr);
+  CHECK(*denominator != 0);
+  int64_t divisor = std::gcd(std::llabs(*numerator), std::llabs(*denominator));
+  if (divisor > 1) {
+    *numerator /= divisor;
+    *denominator /= divisor;
+  }
+  if (*denominator < 0) {
+    *numerator = -*numerator;
+    *denominator = -*denominator;
+  }
+}
+
+// We normalize affine expressions into:
+//   (constant + sum_i coeff_i * var_i) / denominator
+//
+// Concretely, the canonical forms we want to preserve are:
+//   p
+//   mX
+//   mX + p
+//   mX + nY
+//   (mX) / d
+//   (mX + p) / d
+//   (mX + nY + p) / d
+//
+// This lets us combine equivalent trees into one stable representation, e.g.:
+//   A/2 + A/2     -> A
+//   A/2 + B/2     -> (A + B) / 2
+//   4/8 * A       -> A / 2
+//   A + (B + 1)   -> A + B + 1   (represented via coefficients + constant)
+//
+// The important constraint is that we only canonicalize affine expressions over
+// an integer denominator. Non-affine expressions stay in tree form and only get
+// a minimal local simplification in the fallback path below.
+struct CanonicalAffineExpr {
+  int64_t denominator = 1;
+  int64_t constant = 0;
+  std::map<int, int64_t> coefficients;
+
+  bool IsPureConstant() const { return coefficients.empty(); }
+};
+
+void NormalizeAffine(CanonicalAffineExpr* expr) {
+  CHECK(expr != nullptr);
+  CHECK(expr->denominator != 0);
+  for (auto it = expr->coefficients.begin(); it != expr->coefficients.end();) {
+    if (it->second == 0) {
+      it = expr->coefficients.erase(it);
+    } else {
+      ++it;
+    }
+  }
+  int64_t divisor = std::llabs(expr->constant);
+  for (const auto& [_, coeff] : expr->coefficients) {
+    divisor = std::gcd(divisor, std::llabs(coeff));
+  }
+  divisor = std::gcd(divisor, std::llabs(expr->denominator));
+  if (divisor > 1) {
+    expr->constant /= divisor;
+    expr->denominator /= divisor;
+    for (auto& [_, coeff] : expr->coefficients) {
+      coeff /= divisor;
+    }
+  }
+  if (expr->denominator < 0) {
+    expr->denominator = -expr->denominator;
+    expr->constant = -expr->constant;
+    for (auto& [_, coeff] : expr->coefficients) {
+      coeff = -coeff;
+    }
+  }
+  if (expr->constant == 0 && expr->coefficients.empty()) {
+    expr->denominator = 1;
+  }
+}
+
+CanonicalAffineExpr MakeConstantAffine(int64_t value) {
+  CanonicalAffineExpr expr;
+  expr.constant = value;
+  return expr;
+}
+
+CanonicalAffineExpr MakeVariableAffine(int id) {
+  CanonicalAffineExpr expr;
+  expr.coefficients[id] = 1;
+  return expr;
+}
+
+CanonicalAffineExpr AddAffine(const CanonicalAffineExpr& lhs,
+                              const CanonicalAffineExpr& rhs, int rhs_sign) {
+  CHECK(rhs_sign == 1 || rhs_sign == -1);
+  CanonicalAffineExpr result;
+  int64_t gcd = std::gcd(lhs.denominator, rhs.denominator);
+  int64_t lhs_scale = rhs.denominator / gcd;
+  int64_t rhs_scale = lhs.denominator / gcd;
+  result.denominator = lhs.denominator * lhs_scale;
+  result.constant =
+      lhs.constant * lhs_scale + rhs_sign * rhs.constant * rhs_scale;
+  for (const auto& [id, coeff] : lhs.coefficients) {
+    result.coefficients[id] = coeff * lhs_scale;
+  }
+  for (const auto& [id, coeff] : rhs.coefficients) {
+    result.coefficients[id] += rhs_sign * coeff * rhs_scale;
+  }
+  NormalizeAffine(&result);
+  return result;
+}
+
+CanonicalAffineExpr MultiplyAffineByRational(const CanonicalAffineExpr& expr,
+                                             int64_t numerator,
+                                             int64_t denominator) {
+  CHECK(denominator != 0);
+  CanonicalAffineExpr result = expr;
+  result.constant *= numerator;
+  result.denominator *= denominator;
+  for (auto& [_, coeff] : result.coefficients) {
+    coeff *= numerator;
+  }
+  NormalizeAffine(&result);
+  return result;
+}
+
+std::optional<CanonicalAffineExpr> ToCanonicalAffine(const DynExpr* expr) {
+  CHECK(expr != nullptr);
+  switch (expr->kind()) {
+    case DExpr::Kind::kUnknown:
+      return std::nullopt;
+    case DExpr::Kind::kConstant:
+      return MakeConstantAffine(static_cast<const Constant*>(expr)->get_val());
+    case DExpr::Kind::kVariable:
+      return MakeVariableAffine(static_cast<const Variable*>(expr)->get_id());
+    case DExpr::Kind::kAdd: {
+      const auto* add = static_cast<const Add*>(expr);
+      auto lhs = ToCanonicalAffine(add->get_lhs());
+      auto rhs = ToCanonicalAffine(add->get_rhs());
+      if (!lhs.has_value() || !rhs.has_value()) return std::nullopt;
+      return AddAffine(*lhs, *rhs, /*rhs_sign=*/1);
+    }
+    case DExpr::Kind::kSub: {
+      const auto* sub = static_cast<const Sub*>(expr);
+      auto lhs = ToCanonicalAffine(sub->get_lhs());
+      auto rhs = ToCanonicalAffine(sub->get_rhs());
+      if (!lhs.has_value() || !rhs.has_value()) return std::nullopt;
+      return AddAffine(*lhs, *rhs, /*rhs_sign=*/-1);
+    }
+    case DExpr::Kind::kMul: {
+      const auto* mul = static_cast<const Mul*>(expr);
+      auto lhs = ToCanonicalAffine(mul->get_lhs());
+      auto rhs = ToCanonicalAffine(mul->get_rhs());
+      if (!lhs.has_value() || !rhs.has_value()) return std::nullopt;
+      if (lhs->IsPureConstant()) {
+        return MultiplyAffineByRational(*rhs, lhs->constant, lhs->denominator);
+      }
+      if (rhs->IsPureConstant()) {
+        return MultiplyAffineByRational(*lhs, rhs->constant, rhs->denominator);
+      }
+      return std::nullopt;
+    }
+    case DExpr::Kind::kDiv: {
+      const auto* div = static_cast<const Div*>(expr);
+      auto lhs = ToCanonicalAffine(div->get_lhs());
+      auto rhs = ToCanonicalAffine(div->get_rhs());
+      if (!lhs.has_value() || !rhs.has_value() || !rhs->IsPureConstant() ||
+          rhs->constant == 0) {
+        return std::nullopt;
+      }
+      return MultiplyAffineByRational(*lhs, rhs->denominator, rhs->constant);
+    }
+    case DExpr::Kind::kMax:
+    case DExpr::Kind::kGt:
+    case DExpr::Kind::kSelect:
+      return std::nullopt;
+    default:
+      return std::nullopt;
+  }
+  return std::nullopt;
+}
+
+bool IsNonNegativeForPositiveVariables(const DynExpr* expr) {
+  auto affine = ToCanonicalAffine(expr);
+  if (!affine.has_value()) return false;
+
+  // Dynamic dimension variables are strictly positive. An affine expression
+  // has a finite lower bound only when every variable coefficient is
+  // non-negative; evaluate that bound with each variable set to one.
+  __int128 lower_bound = affine->constant;
+  for (const auto& [_, coefficient] : affine->coefficients) {
+    if (coefficient < 0) return false;
+    lower_bound += coefficient;
+  }
+  return lower_bound >= 0;
+}
+
+std::unique_ptr<DynExpr> BuildScaledVariableTerm(int id, int64_t coefficient) {
+  CHECK(coefficient != 0);
+  if (coefficient == 1) {
+    return std::make_unique<Variable>(id);
+  }
+  return std::make_unique<Mul>(DynExpr::_(coefficient), DynExpr::V(id));
+}
+
+std::unique_ptr<DynExpr> BuildAffineNumerator(const CanonicalAffineExpr& expr) {
+  std::unique_ptr<DynExpr> result;
+  for (const auto& [id, coeff] : expr.coefficients) {
+    auto term = BuildScaledVariableTerm(id, coeff);
+    if (result == nullptr) {
+      result = std::move(term);
+    } else {
+      result = std::make_unique<Add>(result.release(), term.release());
+    }
+  }
+  if (expr.constant != 0 || result == nullptr) {
+    if (result == nullptr) {
+      result = std::make_unique<Constant>(expr.constant);
+    } else if (expr.constant > 0) {
+      result = std::make_unique<Add>(result.release(),
+                                     DynExpr::_(expr.constant));
+    } else {
+      result = std::make_unique<Sub>(result.release(),
+                                     DynExpr::_(-expr.constant));
+    }
+  }
+  return result;
+}
+
+std::unique_ptr<DynExpr> BuildCanonicalExpr(const CanonicalAffineExpr& expr) {
+  CanonicalAffineExpr normalized = expr;
+  NormalizeAffine(&normalized);
+  auto numerator = BuildAffineNumerator(normalized);
+  if (normalized.denominator == 1) {
+    return numerator;
+  }
+  return std::make_unique<Div>(numerator.release(),
+                               DynExpr::_(normalized.denominator));
+}
+
+std::unique_ptr<DynExpr> SimplifyFallback(const DynExpr* expr) {
+  switch (expr->kind()) {
+    case DExpr::Kind::kUnknown:
+    case DExpr::Kind::kConstant:
+    case DExpr::Kind::kVariable:
+      return expr->clone();
+    case DExpr::Kind::kAdd: {
+      const auto* add = static_cast<const Add*>(expr);
+      auto lhs = std::unique_ptr<DynExpr>(add->get_lhs()->s());
+      auto rhs = std::unique_ptr<DynExpr>(add->get_rhs()->s());
+      if (lhs->kind() == DExpr::Kind::kUnknown ||
+          rhs->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      Constant* l = AsConstant(lhs.get());
+      Constant* r = AsConstant(rhs.get());
+      if (l && r) return std::make_unique<Constant>(l->get_val() + r->get_val());
+      if (l && l->get_val() == 0) return rhs;
+      if (r && r->get_val() == 0) return lhs;
+      return std::make_unique<Add>(lhs.release(), rhs.release());
+    }
+    case DExpr::Kind::kSub: {
+      const auto* sub = static_cast<const Sub*>(expr);
+      auto lhs = std::unique_ptr<DynExpr>(sub->get_lhs()->s());
+      auto rhs = std::unique_ptr<DynExpr>(sub->get_rhs()->s());
+      if (lhs->kind() == DExpr::Kind::kUnknown ||
+          rhs->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      Constant* l = AsConstant(lhs.get());
+      Constant* r = AsConstant(rhs.get());
+      if (l && r) return std::make_unique<Constant>(l->get_val() - r->get_val());
+      if (r && r->get_val() == 0) return lhs;
+      if (*lhs == *rhs) return std::make_unique<Constant>(0);
+      return std::make_unique<Sub>(lhs.release(), rhs.release());
+    }
+    case DExpr::Kind::kMul: {
+      const auto* mul = static_cast<const Mul*>(expr);
+      auto lhs = std::unique_ptr<DynExpr>(mul->get_lhs()->s());
+      auto rhs = std::unique_ptr<DynExpr>(mul->get_rhs()->s());
+      if (lhs->kind() == DExpr::Kind::kUnknown ||
+          rhs->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      Constant* l = AsConstant(lhs.get());
+      Constant* r = AsConstant(rhs.get());
+      if (l && r) return std::make_unique<Constant>(l->get_val() * r->get_val());
+      if ((l && l->get_val() == 0) || (r && r->get_val() == 0)) {
+        return std::make_unique<Constant>(0);
+      }
+      if (l && l->get_val() == 1) return rhs;
+      if (r && r->get_val() == 1) return lhs;
+      if (r != nullptr) std::swap(lhs, rhs);
+      return std::make_unique<Mul>(lhs.release(), rhs.release());
+    }
+    case DExpr::Kind::kDiv: {
+      const auto* div = static_cast<const Div*>(expr);
+      auto lhs = std::unique_ptr<DynExpr>(div->get_lhs()->s());
+      auto rhs = std::unique_ptr<DynExpr>(div->get_rhs()->s());
+      if (lhs->kind() == DExpr::Kind::kUnknown ||
+          rhs->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      Constant* l = AsConstant(lhs.get());
+      Constant* r = AsConstant(rhs.get());
+      if (*lhs == *rhs) {
+        return std::make_unique<Constant>(1);
+      }
+      if (l && l->get_val() == 0 && r && r->get_val() != 0) {
+        return std::make_unique<Constant>(0);
+      }
+      if (r && r->get_val() == 1) return lhs;
+      if (lhs->kind() == DExpr::Kind::kMul) {
+        auto* mul = static_cast<Mul*>(lhs.get());
+        auto lhs_l = std::unique_ptr<DynExpr>(mul->get_lhs()->s());
+        auto lhs_r = std::unique_ptr<DynExpr>(mul->get_rhs()->s());
+        if (*lhs_l == *rhs) {
+          return lhs_r;
+        }
+        if (*lhs_r == *rhs) {
+          return lhs_l;
+        }
+      }
+      if (l && r && r->get_val() != 0) {
+        int64_t numerator = l->get_val();
+        int64_t denominator = r->get_val();
+        NormalizeFraction(&numerator, &denominator);
+        if (denominator == 1) return std::make_unique<Constant>(numerator);
+        return std::make_unique<Div>(DynExpr::_(numerator),
+                                     DynExpr::_(denominator));
+      }
+      return std::make_unique<Div>(lhs.release(), rhs.release());
+    }
+    case DExpr::Kind::kMax: {
+      const auto* max = static_cast<const MaxExpr*>(expr);
+      auto lhs = std::unique_ptr<DynExpr>(max->get_lhs()->s());
+      auto rhs = std::unique_ptr<DynExpr>(max->get_rhs()->s());
+      if (lhs->kind() == DExpr::Kind::kUnknown ||
+          rhs->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      Constant* l = AsConstant(lhs.get());
+      Constant* r = AsConstant(rhs.get());
+      if (l && r) {
+        return std::make_unique<Constant>(
+            std::max(l->get_val(), r->get_val()));
+      }
+      if (l && l->get_val() == 0 &&
+          IsNonNegativeForPositiveVariables(rhs.get())) {
+        return rhs;
+      }
+      if (r && r->get_val() == 0 &&
+          IsNonNegativeForPositiveVariables(lhs.get())) {
+        return lhs;
+      }
+      if (*lhs == *rhs) return lhs;
+      return std::make_unique<MaxExpr>(lhs.release(), rhs.release());
+    }
+    case DExpr::Kind::kGt: {
+      const auto* gt = static_cast<const GtExpr*>(expr);
+      auto lhs = std::unique_ptr<DynExpr>(gt->get_lhs()->s());
+      auto rhs = std::unique_ptr<DynExpr>(gt->get_rhs()->s());
+      if (lhs->kind() == DExpr::Kind::kUnknown ||
+          rhs->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      if (lhs->is_constant() && rhs->is_constant()) {
+        return std::make_unique<Constant>(lhs->get_val() > rhs->get_val());
+      }
+      if (*lhs == *rhs) return std::make_unique<Constant>(0);
+      return std::make_unique<GtExpr>(lhs.release(), rhs.release());
+    }
+    case DExpr::Kind::kSelect: {
+      const auto* select = static_cast<const SelectExpr*>(expr);
+      auto pred = std::unique_ptr<DynExpr>(select->get_pred()->s());
+      auto on_true = std::unique_ptr<DynExpr>(select->get_on_true()->s());
+      auto on_false = std::unique_ptr<DynExpr>(select->get_on_false()->s());
+      if (pred->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      if (pred->is_constant()) {
+        return pred->get_val() != 0 ? std::move(on_true) : std::move(on_false);
+      }
+      if (on_true->kind() == DExpr::Kind::kUnknown ||
+          on_false->kind() == DExpr::Kind::kUnknown) {
+        return std::make_unique<UnknownExpr>();
+      }
+      if (*on_true == *on_false) return on_true;
+      return std::make_unique<SelectExpr>(pred.release(), on_true.release(),
+                                           on_false.release());
+    }
+    default:
+      return expr->clone();
+  }
+}
+
+std::unique_ptr<DynExpr> SimplifyCanonical(const DynExpr* expr) {
+  if (expr->kind() == DExpr::Kind::kUnknown) {
+    return std::make_unique<UnknownExpr>();
+  }
+  if (auto canonical = ToCanonicalAffine(expr); canonical.has_value()) {
+    if (canonical->IsPureConstant()) {
+      CHECK_NE(canonical->denominator, 0);
+      return std::make_unique<Constant>(canonical->constant /
+                                        canonical->denominator);
+    }
+    return BuildCanonicalExpr(*canonical);
+  }
+  return SimplifyFallback(expr);
+}
+
+}  // namespace
+
+DExpr DExpr::find_smallest_subexpression_covering_all_variables() const {
+  CHECK(expr_ != nullptr);
+  const std::set<int> ids = expr_->get_all_ids();
+  CHECK(!ids.empty());
+  DynExpr* result = FindSmallestCoveringSubexpression(expr_.get());
+  CHECK(result != nullptr);
+  return DExpr::Adopt(result->clone().release());
+}
+
+DExpr DExpr::replace_subexpression(const DExpr& target,
+                                   const DExpr& replacement) const {
+  if (expr_ == nullptr) {
+    return DExpr();
+  }
+  return DExpr(ReplaceSubexpression(expr_.get(), target.get(),
+                                    replacement.get()));
+}
+
+const DExpr& Shape::MissingExpression() {
+  static const DExpr missing = DExpr::Unknown(kMissingExpressionSentinel);
+  return missing;
+}
+
+DynExpr* operator*(DynExpr& lhs, DynExpr& rhs) {
+  return new Mul(lhs.clone().release(), rhs.clone().release());
+}
+DynExpr* operator*(int64_t k, DynExpr& rhs) {
+  return new Mul(DynExpr::_(k), rhs.clone().release());
+}
+DynExpr* operator*(DynExpr& lhs, int64_t k) {
+  return new Mul(lhs.clone().release(), DynExpr::_(k));
+}
+DynExpr* operator/(DynExpr& lhs, DynExpr& rhs) {
+  return new Div(lhs.clone().release(), rhs.clone().release());
+}
+DynExpr* operator/(DynExpr& lhs, int64_t d) {
+  return new Div(lhs.clone().release(), DynExpr::_(d));
+}
+DynExpr* operator+(DynExpr& lhs, DynExpr& rhs) {
+  return new Add(lhs.clone().release(), rhs.clone().release());
+}
+DynExpr* operator+(DynExpr& lhs, int64_t d) {
+  return new Add(lhs.clone().release(), DynExpr::_(d));
+}
+DynExpr* operator-(DynExpr& lhs, DynExpr& rhs) {
+  return new Sub(lhs.clone().release(), rhs.clone().release());
+}
+DynExpr* operator-(DynExpr& lhs, int64_t d) {
+  return new Sub(lhs.clone().release(), DynExpr::_(d));
+}
+bool operator==(DynExpr& lhs, DynExpr& rhs) {
+  return DynExpr::equal(&lhs, &rhs);
+}
+bool operator==(DynExpr& lhs, int64_t d) {
+  auto rhs = std::unique_ptr<DynExpr>(DynExpr::_(d));
+  return DynExpr::equal(&lhs, rhs.get());
+}
+bool operator<(DynExpr& lhs, int64_t d) {
+  return lhs.is_constant() && lhs.get_val() < d;
+}
+
+DExpr DExpr::Max(const DExpr& lhs, const DExpr& rhs) {
+  return Adopt(new xla::MaxExpr(lhs.clone().release(), rhs.clone().release()));
+}
+
+DExpr DExpr::Gt(const DExpr& lhs, const DExpr& rhs) {
+  return Adopt(new xla::GtExpr(lhs.clone().release(), rhs.clone().release()));
+}
+
+DExpr DExpr::Select(const DExpr& pred, const DExpr& on_true,
+                    const DExpr& on_false) {
+  return Adopt(new xla::SelectExpr(pred.clone().release(),
+                                    on_true.clone().release(),
+                                    on_false.clone().release()));
+}
+
+bool DynExpr::equal(DynExpr* expr1, DynExpr* expr2) {
+  auto e1 = std::unique_ptr<DynExpr>(expr1->s());
+  auto e2 = std::unique_ptr<DynExpr>(expr2->s());
+  if (e1 == nullptr || e2 == nullptr) return false;
+  auto a1 = ToCanonicalAffine(e1.get());
+  auto a2 = ToCanonicalAffine(e2.get());
+  if (a1.has_value() && a2.has_value()) {
+    return a1->denominator == a2->denominator &&
+           a1->constant == a2->constant &&
+           a1->coefficients == a2->coefficients;
+  }
+  if (e1->kind() == DExpr::Kind::kConstant &&
+      e2->kind() == DExpr::Kind::kConstant) {
+    return static_cast<Constant*>(e1.get())->get_val() ==
+           static_cast<Constant*>(e2.get())->get_val();
+  }
+  if (e1->kind() == DExpr::Kind::kVariable &&
+      e2->kind() == DExpr::Kind::kVariable) {
+    return static_cast<Variable*>(e1.get())->get_id() ==
+           static_cast<Variable*>(e2.get())->get_id();
+  }
+  if (e1->kind() == DExpr::Kind::kUnknown &&
+      e2->kind() == DExpr::Kind::kUnknown) {
+    int lhs_id = static_cast<UnknownExpr*>(e1.get())->get_id();
+    int rhs_id = static_cast<UnknownExpr*>(e2.get())->get_id();
+    return lhs_id != 0 && lhs_id == rhs_id;
+  }
+  if (e1->kind() == DExpr::Kind::kMul && e2->kind() == DExpr::Kind::kMul) {
+    auto* ab = static_cast<Mul*>(e1.get());
+    auto* cd = static_cast<Mul*>(e2.get());
+    auto a = ab->get_lhs();
+    auto b = ab->get_rhs();
+    auto c = cd->get_lhs();
+    auto d = cd->get_rhs();
+    return (*a == *c && *b == *d) || (*a == *d && *b == *c);
+  }
+  if (e1->kind() == DExpr::Kind::kDiv && e2->kind() == DExpr::Kind::kDiv) {
+    auto* ab = static_cast<Div*>(e1.get());
+    auto* cd = static_cast<Div*>(e2.get());
+    auto a = ab->get_lhs();
+    auto b = ab->get_rhs();
+    auto c = cd->get_lhs();
+    auto d = cd->get_rhs();
+    return *a == *c && *b == *d;
+  }
+  if (e1->kind() == DExpr::Kind::kAdd && e2->kind() == DExpr::Kind::kAdd) {
+    auto* ab = static_cast<Add*>(e1.get());
+    auto* cd = static_cast<Add*>(e2.get());
+    auto a = ab->get_lhs();
+    auto b = ab->get_rhs();
+    auto c = cd->get_lhs();
+    auto d = cd->get_rhs();
+    return (*a == *c && *b == *d) || (*a == *d && *b == *c);
+  }
+  if (e1->kind() == DExpr::Kind::kSub && e2->kind() == DExpr::Kind::kSub) {
+    auto* ab = static_cast<Sub*>(e1.get());
+    auto* cd = static_cast<Sub*>(e2.get());
+    auto* a = ab->get_lhs();
+    auto* b = ab->get_rhs();
+    auto* c = cd->get_lhs();
+    auto* d = cd->get_rhs();
+    return *a == *c && *b == *d;
+  }
+  if (e1->kind() == DExpr::Kind::kMax && e2->kind() == DExpr::Kind::kMax) {
+    auto* ab = static_cast<MaxExpr*>(e1.get());
+    auto* cd = static_cast<MaxExpr*>(e2.get());
+    auto* a = ab->get_lhs();
+    auto* b = ab->get_rhs();
+    auto* c = cd->get_lhs();
+    auto* d = cd->get_rhs();
+    return (*a == *c && *b == *d) || (*a == *d && *b == *c);
+  }
+  if (e1->kind() == DExpr::Kind::kGt && e2->kind() == DExpr::Kind::kGt) {
+    auto* lhs = static_cast<GtExpr*>(e1.get());
+    auto* rhs = static_cast<GtExpr*>(e2.get());
+    return *lhs->get_lhs() == *rhs->get_lhs() &&
+           *lhs->get_rhs() == *rhs->get_rhs();
+  }
+  if (e1->kind() == DExpr::Kind::kSelect &&
+      e2->kind() == DExpr::Kind::kSelect) {
+    auto* lhs = static_cast<SelectExpr*>(e1.get());
+    auto* rhs = static_cast<SelectExpr*>(e2.get());
+    return *lhs->get_pred() == *rhs->get_pred() &&
+           *lhs->get_on_true() == *rhs->get_on_true() &&
+           *lhs->get_on_false() == *rhs->get_on_false();
+  }
+  return false;
+}
+
+DynExpr* Constant::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* Variable::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* Mul::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* Add::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* Sub::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* Div::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* MaxExpr::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* GtExpr::s() { return SimplifyCanonical(this).release(); }
+
+DynExpr* SelectExpr::s() { return SimplifyCanonical(this).release(); }
+
+std::ostream& operator<<(std::ostream& os, DynExpr* expr) {
+  auto simplified = std::unique_ptr<DynExpr>(expr->s());
+  StringPrinter printer;
+  simplified->print(&printer);
+  os << std::move(printer).ToString();
+  return os;
+}
+
+DynExpr* DynExpr::zero = new Constant(0);
+DynExpr* DynExpr::one = new Constant(1);
+
+}  // namespace xla
diff --git a/third_party/xla/xla/shape_expr.h b/third_party/xla/xla/shape_expr.h
new file mode 100644
index 00000000000000..49b564fbb92a57
--- /dev/null
+++ b/third_party/xla/xla/shape_expr.h
@@ -0,0 +1,821 @@
+/* Copyright 2018 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef XLA_SHAPE_EXPR_H_
+#define XLA_SHAPE_EXPR_H_
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <ostream>
+#include <set>
+#include <utility>
+#include <vector>
+
+#include "absl/hash/hash.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/types/span.h"
+#include "xla/printer.h"
+#include "xla/xla_data.pb.h"
+
+namespace xla {
+
+// Reserved sentinel for "missing expression". Keep this outside the normal
+// expression id space so it cannot be confused with a real UnknownExpr id.
+inline constexpr int kMissingExpressionSentinel = -1000001;
+inline constexpr int64_t kUnknownContentSentinel = -444;
+
+enum class DExprKind {
+  kUnknown,
+  kConstant,
+  kVariable,
+  kAdd,
+  kSub,
+  kMul,
+  kDiv,
+  kMax,
+  kGt,
+  kSelect,
+};
+
+class DynExpr {
+ public:
+  virtual ~DynExpr() = default;
+  virtual std::unique_ptr<DynExpr> clone() const = 0;
+  virtual DExprKind kind() const = 0;
+  virtual void print(xla::Printer* printer) const = 0;
+  virtual void to_proto(xla::ExpressionProto* proto) const = 0;
+  virtual bool is_constant() const = 0;
+  virtual int64_t get_val() const { return -1; }
+  virtual DynExpr* s() = 0; // simplify
+  virtual DynExpr* substitute(int id, DynExpr* v) = 0;
+  virtual std::set<int> get_all_ids() = 0;
+  virtual std::optional<int64_t> solve(int64_t x) = 0;
+
+  bool is_dynamic() { return !is_constant(); }
+
+  static DynExpr* zero;
+  static DynExpr* one;
+  static DynExpr* _(int64_t val);
+  static DynExpr* V(int var_id);
+  static DynExpr* _s(DynExpr* expr);
+  static bool equal(DynExpr* expr1, DynExpr* expr2);
+
+  friend std::ostream& operator<<(std::ostream& os, DynExpr* expr);
+};
+
+class DExpr {
+ public:
+  using Kind = DExprKind;
+
+  DExpr() = default;
+  explicit DExpr(std::unique_ptr<DynExpr> expr) : expr_(std::move(expr)) {}
+
+  DExpr(const DExpr& other) {
+    if (other.expr_ != nullptr) {
+      expr_ = other.expr_->clone();
+    }
+  }
+  DExpr& operator=(const DExpr& other) {
+    if (this == &other) return *this;
+    expr_.reset();
+    if (other.expr_ != nullptr) {
+      expr_ = other.expr_->clone();
+    }
+    return *this;
+  }
+
+  DExpr(DExpr&&) noexcept = default;
+  DExpr& operator=(DExpr&&) noexcept = default;
+
+  static DExpr Unknown(int id = 0);
+  static DExpr Adopt(DynExpr* expr) { return DExpr(std::unique_ptr<DynExpr>(expr)); }
+  static DExpr Const(int64_t value) { return Adopt(DynExpr::_(value)); }
+  static DExpr Var(int var_id) { return Adopt(DynExpr::V(var_id)); }
+  static DExpr Max(const DExpr& lhs, const DExpr& rhs);
+  static DExpr Gt(const DExpr& lhs, const DExpr& rhs);
+  static DExpr Select(const DExpr& pred, const DExpr& on_true,
+                      const DExpr& on_false);
+  bool is_unknown() const {
+    return expr_ != nullptr && expr_->kind() == DExprKind::kUnknown;
+  }
+  Kind kind() const {
+    CHECK(expr_ != nullptr) << "Attempted to access empty DExpr";
+    return expr_->kind();
+  }
+
+  DynExpr* get() const {
+    CHECK(expr_ != nullptr) << "Attempted to access empty DExpr";
+    return expr_.get();
+  }
+  DynExpr& operator*() const {
+    CHECK(expr_ != nullptr) << "Attempted to dereference empty DExpr";
+    return *expr_;
+  }
+  DynExpr* operator->() const {
+    CHECK(expr_ != nullptr) << "Attempted to access empty DExpr";
+    return expr_.get();
+  }
+  operator DynExpr*() const { return get(); }
+  explicit operator bool() const { return expr_ != nullptr && !is_unknown(); }
+
+  std::unique_ptr<DynExpr> clone() const {
+    if (expr_ == nullptr) {
+      return nullptr;
+    }
+    return expr_->clone();
+  }
+  DynExpr* release() { return expr_.release(); }
+
+  DExpr simplify() const {
+    return expr_ == nullptr ? DExpr() : Adopt(expr_->s());
+  }
+  void to_proto(xla::ExpressionProto* proto) const {
+    CHECK(expr_ != nullptr) << "Attempted to serialize empty DExpr";
+    expr_->to_proto(proto);
+  }
+  DExpr substitute(int id, const DExpr& value) const {
+    return expr_ == nullptr ? DExpr() : Adopt(expr_->substitute(id, value.get()));
+  }
+  // Returns the smallest subtree containing every variable in this expression.
+  // The expression must contain at least one variable.
+  DExpr find_smallest_subexpression_covering_all_variables() const;
+  // Replaces every subtree equivalent to `target` with `replacement`.
+  DExpr replace_subexpression(const DExpr& target,
+                              const DExpr& replacement) const;
+
+  template <typename H>
+  friend H AbslHashValue(H h, const DExpr& expr) {
+    xla::ExpressionProto proto;
+    if (expr.expr_ != nullptr) {
+      expr.expr_->to_proto(&proto);
+    }
+    return H::combine(std::move(h), proto.SerializeAsString());
+  }
+
+ private:
+  std::unique_ptr<DynExpr> expr_;
+};
+
+class UnknownExpr : public DynExpr {
+  int id_;
+
+ public:
+  explicit UnknownExpr(int id = 0) : id_(id) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<UnknownExpr>(id_);
+  }
+  DExprKind kind() const override { return DExprKind::kUnknown; }
+  void print(xla::Printer* printer) const override {
+    if (id_ == kMissingExpressionSentinel || id_ == kUnknownContentSentinel) {
+      printer->Append("_");
+      return;
+    }
+    printer->Append("?");
+    if (id_ != 0) {
+      printer->Append(id_);
+    }
+  }
+  void to_proto(xla::ExpressionProto* proto) const override {
+    (void)proto;
+  }
+  bool is_constant() const override { return true; }
+  int get_id() const { return id_; }
+  DynExpr* substitute(int id, DynExpr* v) override {
+    (void)id;
+    (void)v;
+    return clone().release();
+  }
+  std::set<int> get_all_ids() override { return {}; }
+  std::optional<int64_t> solve(int64_t) override {
+    return std::nullopt;
+  }
+  DynExpr* s() override { return clone().release(); }
+};
+
+inline DExpr DExpr::Unknown(int id) {
+  return DExpr(std::unique_ptr<DynExpr>(new xla::UnknownExpr(id)));
+}
+
+// constant i
+class Constant : public DynExpr {
+  int64_t value;
+
+ public:
+  explicit Constant(int64_t v) : value(v) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<Constant>(value);
+  }
+  DExprKind kind() const override { return DExprKind::kConstant; }
+  void print(xla::Printer* printer) const override {
+    if (value < 0) {
+      printer->Append("(");
+    }
+    printer->Append(value);
+    if (value < 0) {
+      printer->Append(")");
+    }
+  }
+  void to_proto(xla::ExpressionProto* proto) const override {
+    proto->set_constant_value(value);
+  }
+  bool is_constant() const override { return true; }
+  int64_t get_val() const override { return value; }
+  DynExpr* substitute(int id, DynExpr* v) { return clone().release(); }
+  std::set<int> get_all_ids() { return {}; }
+  std::optional<int64_t> solve(int64_t x) { return std::nullopt; }
+  DynExpr* s() override;
+};
+
+// Root variables represent strictly positive dynamic dimensions. Potentially
+// signed values must be represented by expressions derived from these roots.
+class Variable : public DynExpr {
+  int id;
+
+ public:
+  explicit Variable(int identifier) : id(identifier) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<Variable>(id);
+  }
+  DExprKind kind() const override { return DExprKind::kVariable; }
+  void print(xla::Printer* printer) const override {
+    if (id >= 1 && id <= 26) {
+      char letter = 'A' + (id - 1);
+      printer->Append(std::string(1, letter));
+      return;
+    }
+    printer->Append("V(");
+    printer->Append(id);
+    printer->Append(")");
+  }
+  void to_proto(xla::ExpressionProto* proto) const override {
+    proto->set_variable_id(id);
+  }
+  bool is_constant() const override { return false; }
+  int get_id() const { return id; }
+  DynExpr* substitute(int id, DynExpr* v) {
+    return get_id() == id ? v->clone().release() : clone().release();
+  }
+  std::set<int> get_all_ids() { return {get_id()}; }
+  std::optional<int64_t> solve(int64_t x) { return x; }
+  DynExpr* s() override;
+};
+
+// exp = exp + exp
+class Add : public DynExpr {
+  std::unique_ptr<DynExpr> lhs;
+  std::unique_ptr<DynExpr> rhs;
+
+ public:
+  Add(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<Add>(lhs->clone().release(), rhs->clone().release());
+  }
+  DExprKind kind() const override { return DExprKind::kAdd; }
+  void print(xla::Printer* printer) const override {
+    printer->Append("(");
+    lhs->print(printer);
+    printer->Append(" + ");
+    rhs->print(printer);
+    printer->Append(")");
+  }
+
+  DynExpr* get_lhs() const { return lhs.get(); }
+  DynExpr* get_rhs() const { return rhs.get(); }
+
+  void to_proto(xla::ExpressionProto* proto) const override {
+    auto* add_msg = proto->mutable_add_node();
+    lhs->to_proto(add_msg->mutable_lhs());
+    rhs->to_proto(add_msg->mutable_rhs());
+  }
+
+  bool is_constant() const override {
+    return lhs->is_constant() && rhs->is_constant();
+  }
+
+  int64_t get_val() const override { return lhs->get_val() + rhs->get_val(); }
+
+  DynExpr* substitute(int id, DynExpr* v) {
+    return new Add(lhs->substitute(id, v), rhs->substitute(id, v));
+  }
+
+  std::set<int> get_all_ids() {
+    auto s = lhs->get_all_ids();
+    s.merge(rhs->get_all_ids());
+    return s;
+  }
+
+  std::optional<int64_t> solve(int64_t x) {
+    // Cannot solve if both lhs and rhs are dynamic...
+    if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt;
+    if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) {
+      // (A + c) = x <=> A = x - c => solve A = y with y = x - c
+      return lhs->solve(x - rhs->get_val());
+    }
+    if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) {
+      // (c + A) = x <=> A = x - c => solve A = y with y = x - c
+      return rhs->solve(x - lhs->get_val());
+    }
+    // No solution
+    return std::nullopt;
+  }
+
+  DynExpr* s() override;
+
+  ~Add() override = default;
+};
+
+// exp = exp - exp
+class Sub : public DynExpr {
+  std::unique_ptr<DynExpr> lhs;
+  std::unique_ptr<DynExpr> rhs;
+
+ public:
+  Sub(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<Sub>(lhs->clone().release(), rhs->clone().release());
+  }
+  DExprKind kind() const override { return DExprKind::kSub; }
+  void print(xla::Printer* printer) const override {
+    printer->Append("(");
+    lhs->print(printer);
+    printer->Append(" - ");
+    rhs->print(printer);
+    printer->Append(")");
+  }
+
+  DynExpr* get_lhs() const { return lhs.get(); }
+  DynExpr* get_rhs() const { return rhs.get(); }
+
+  void to_proto(xla::ExpressionProto* proto) const override {
+    auto* sub_msg = proto->mutable_sub_node();
+    lhs->to_proto(sub_msg->mutable_lhs());
+    rhs->to_proto(sub_msg->mutable_rhs());
+  }
+
+  bool is_constant() const override {
+    return lhs->is_constant() && rhs->is_constant();
+  }
+
+  int64_t get_val() const override { return lhs->get_val() - rhs->get_val(); }
+
+  DynExpr* substitute(int id, DynExpr* v) {
+    return new Sub(lhs->substitute(id, v), rhs->substitute(id, v));
+  }
+
+  std::set<int> get_all_ids() {
+    auto s = lhs->get_all_ids();
+    s.merge(rhs->get_all_ids());
+    return s;
+  }
+
+  std::optional<int64_t> solve(int64_t x) {
+    // Cannot solve if both lhs and rhs are dynamic...
+    if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt;
+    if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) {
+      // (A - c) = x <=> A = x + c => solve A = y with y = x + c
+      return lhs->solve(x + rhs->get_val());
+    }
+    if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) {
+      // (c - A) = x <=> A = c - x => solve A = y with y = c - x
+      return rhs->solve(lhs->get_val() - x);
+    }
+    // No solution
+    return std::nullopt;
+  }
+
+  DynExpr* s() override;
+
+  ~Sub() override = default;
+};
+
+// exp = exp * exp
+class Mul : public DynExpr {
+  std::unique_ptr<DynExpr> lhs;
+  std::unique_ptr<DynExpr> rhs;
+
+ public:
+  Mul(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<Mul>(lhs->clone().release(), rhs->clone().release());
+  }
+  DExprKind kind() const override { return DExprKind::kMul; }
+  void print(xla::Printer* printer) const override {
+    printer->Append("(");
+    lhs->print(printer);
+    printer->Append(" * ");
+    rhs->print(printer);
+    printer->Append(")");
+  }
+
+  DynExpr* get_lhs() const { return lhs.get(); }
+  DynExpr* get_rhs() const { return rhs.get(); }
+
+  void to_proto(xla::ExpressionProto* proto) const override {
+    auto* mul_msg = proto->mutable_mul_node();
+    lhs->to_proto(mul_msg->mutable_lhs());
+    rhs->to_proto(mul_msg->mutable_rhs());
+  }
+
+  bool is_constant() const override {
+    return lhs->is_constant() && rhs->is_constant();
+  }
+
+  int64_t get_val() const override { return lhs->get_val() * rhs->get_val(); }
+
+  DynExpr* substitute(int id, DynExpr* v) {
+    return new Mul(lhs->substitute(id, v), rhs->substitute(id, v));
+  }
+
+  std::set<int> get_all_ids() {
+    auto s = lhs->get_all_ids();
+    s.merge(rhs->get_all_ids());
+    return s;
+  }
+
+  std::optional<int64_t> solve(int64_t x) {
+    // Cannot solve if both lhs and rhs are dynamic...
+    if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt;
+    if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) {
+      // (A * c) = x <=> A = x / c => solve A = y with y = x / c
+      int64_t c = rhs->get_val();
+      if (c == 0) return x == 0 ? lhs->solve(0) : std::nullopt;
+      if (x % c != 0) return std::nullopt;
+      return lhs->solve(x / c);
+    }
+    if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) {
+      // (c * A) = x <=> A = x / c => solve A = y with y = x / c
+      int64_t c = lhs->get_val();
+      if (c == 0) return x == 0 ? rhs->solve(0) : std::nullopt;
+      if (x % c != 0) return std::nullopt;
+      return rhs->solve(x / c);
+    }
+    // No solution
+    return std::nullopt;
+  }
+
+  DynExpr* s() override;
+
+  ~Mul() override = default;
+};
+
+// expr / expr
+class Div : public DynExpr {
+  std::unique_ptr<DynExpr> lhs;
+  std::unique_ptr<DynExpr> rhs;
+
+ public:
+  Div(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<Div>(lhs->clone().release(), rhs->clone().release());
+  }
+  DExprKind kind() const override { return DExprKind::kDiv; }
+  void print(xla::Printer* printer) const override {
+    printer->Append("(");
+    lhs->print(printer);
+    printer->Append(" / ");
+    rhs->print(printer);
+    printer->Append(")");
+  }
+
+  DynExpr* get_lhs() const { return lhs.get(); }
+  DynExpr* get_rhs() const { return rhs.get(); }
+
+  void to_proto(xla::ExpressionProto* proto) const override {
+    auto* div_msg = proto->mutable_div_node();
+    lhs->to_proto(div_msg->mutable_lhs());
+    rhs->to_proto(div_msg->mutable_rhs());
+  }
+
+  bool is_constant() const override {
+    return lhs->is_constant() && rhs->is_constant() && rhs->get_val() != 0;
+  }
+
+  int64_t get_val() const override {
+    CHECK(is_constant())
+        << "Attempted to evaluate a non-constant or zero-divisor expression";
+    return lhs->get_val() / rhs->get_val();
+  }
+
+  DynExpr* substitute(int id, DynExpr* v) {
+    return new Div(lhs->substitute(id, v), rhs->substitute(id, v));
+  }
+
+  DynExpr* s() override;
+
+  std::set<int> get_all_ids() {
+    auto s = lhs->get_all_ids();
+    s.merge(rhs->get_all_ids());
+    return s;
+  }
+
+  std::optional<int64_t> solve(int64_t x) {
+    // Cannot solve if both lhs and rhs are dynamic...
+    if (lhs->is_dynamic() && rhs->is_dynamic()) return std::nullopt;
+    if (lhs->get_all_ids().size() == 1 && rhs->is_constant()) {
+      // (A / c) = x <=> A = x * c => solve A = y with y = x * c
+      const int64_t divisor = rhs->get_val();
+      if (divisor == 0) return std::nullopt;
+      return lhs->solve(x * divisor);
+    }
+    if (rhs->get_all_ids().size() == 1 && lhs->is_constant()) {
+      // (c / A) = x <=> A = c / x => solve A = y with y = c / x
+      int64_t c = lhs->get_val();
+      if (x == 0) return std::nullopt;
+      if (c % x != 0) return std::nullopt;
+      return rhs->solve(c / x);
+    }
+    // No solution
+    return std::nullopt;
+  }
+
+  ~Div() override = default;
+};
+
+// max(lhs, rhs)
+class MaxExpr : public DynExpr {
+  std::unique_ptr<DynExpr> lhs;
+  std::unique_ptr<DynExpr> rhs;
+
+ public:
+  MaxExpr(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<MaxExpr>(lhs->clone().release(),
+                                     rhs->clone().release());
+  }
+  DExprKind kind() const override { return DExprKind::kMax; }
+  void print(xla::Printer* printer) const override {
+    printer->Append("max(");
+    lhs->print(printer);
+    printer->Append(", ");
+    rhs->print(printer);
+    printer->Append(")");
+  }
+  void to_proto(xla::ExpressionProto* proto) const override {
+    auto* max_msg = proto->mutable_max_node();
+    lhs->to_proto(max_msg->mutable_lhs());
+    rhs->to_proto(max_msg->mutable_rhs());
+  }
+  bool is_constant() const override {
+    return lhs->is_constant() && rhs->is_constant();
+  }
+  int64_t get_val() const override {
+    return std::max(lhs->get_val(), rhs->get_val());
+  }
+  DynExpr* get_lhs() const { return lhs.get(); }
+  DynExpr* get_rhs() const { return rhs.get(); }
+  DynExpr* substitute(int id, DynExpr* v) override {
+    return new MaxExpr(lhs->substitute(id, v), rhs->substitute(id, v));
+  }
+  std::set<int> get_all_ids() override {
+    auto ids = lhs->get_all_ids();
+    ids.merge(rhs->get_all_ids());
+    return ids;
+  }
+  std::optional<int64_t> solve(int64_t x) override {
+    auto solve_dynamic_branch = [x](DynExpr* dynamic,
+                                    DynExpr* constant)
+        -> std::optional<int64_t> {
+      if (constant->kind() != DExprKind::kConstant ||
+          dynamic->get_all_ids().size() != 1) {
+        return std::nullopt;
+      }
+      const int64_t bound = constant->get_val();
+      // max(expr, bound) selects expr only when the result exceeds bound.
+      // Equality is ambiguous because either branch may have produced it.
+      return x > bound ? dynamic->solve(x) : std::nullopt;
+    };
+
+    if (auto result = solve_dynamic_branch(lhs.get(), rhs.get())) {
+      return result;
+    }
+    if (auto result = solve_dynamic_branch(rhs.get(), lhs.get())) {
+      return result;
+    }
+
+    StringPrinter printer;
+    print(&printer);
+    LOG(WARNING) << "Cannot solve Max dynamic shape expression for value " << x
+                 << ": " << std::move(printer).ToString();
+    return std::nullopt;
+  }
+  DynExpr* s() override;
+};
+
+class GtExpr : public DynExpr {
+  std::unique_ptr<DynExpr> lhs;
+  std::unique_ptr<DynExpr> rhs;
+
+ public:
+  GtExpr(DynExpr* l, DynExpr* r) : lhs(l), rhs(r) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<GtExpr>(lhs->clone().release(),
+                                    rhs->clone().release());
+  }
+  DExprKind kind() const override { return DExprKind::kGt; }
+  void print(xla::Printer* printer) const override {
+    printer->Append("(");
+    lhs->print(printer);
+    printer->Append(" > ");
+    rhs->print(printer);
+    printer->Append(")");
+  }
+  void to_proto(xla::ExpressionProto* proto) const override {
+    auto* gt_msg = proto->mutable_gt_node();
+    lhs->to_proto(gt_msg->mutable_lhs());
+    rhs->to_proto(gt_msg->mutable_rhs());
+  }
+  bool is_constant() const override {
+    return lhs->is_constant() && rhs->is_constant();
+  }
+  int64_t get_val() const override { return lhs->get_val() > rhs->get_val(); }
+  DynExpr* get_lhs() const { return lhs.get(); }
+  DynExpr* get_rhs() const { return rhs.get(); }
+  DynExpr* substitute(int id, DynExpr* v) override {
+    return new GtExpr(lhs->substitute(id, v), rhs->substitute(id, v));
+  }
+  std::set<int> get_all_ids() override {
+    auto ids = lhs->get_all_ids();
+    ids.merge(rhs->get_all_ids());
+    return ids;
+  }
+  std::optional<int64_t> solve(int64_t x) override {
+    StringPrinter printer;
+    print(&printer);
+    LOG(WARNING) << "Cannot solve Gt dynamic shape expression for value " << x
+                 << ": " << std::move(printer).ToString();
+    return std::nullopt;
+  }
+  DynExpr* s() override;
+};
+
+class SelectExpr : public DynExpr {
+  std::unique_ptr<DynExpr> pred;
+  std::unique_ptr<DynExpr> on_true;
+  std::unique_ptr<DynExpr> on_false;
+
+ public:
+  SelectExpr(DynExpr* p, DynExpr* t, DynExpr* f)
+      : pred(p), on_true(t), on_false(f) {}
+  std::unique_ptr<DynExpr> clone() const override {
+    return std::make_unique<SelectExpr>(pred->clone().release(),
+                                        on_true->clone().release(),
+                                        on_false->clone().release());
+  }
+  DExprKind kind() const override { return DExprKind::kSelect; }
+  void print(xla::Printer* printer) const override {
+    printer->Append("select(");
+    pred->print(printer);
+    printer->Append(", ");
+    on_true->print(printer);
+    printer->Append(", ");
+    on_false->print(printer);
+    printer->Append(")");
+  }
+  void to_proto(xla::ExpressionProto* proto) const override {
+    auto* select_msg = proto->mutable_select_node();
+    pred->to_proto(select_msg->mutable_pred());
+    on_true->to_proto(select_msg->mutable_on_true());
+    on_false->to_proto(select_msg->mutable_on_false());
+  }
+  bool is_constant() const override {
+    return pred->is_constant() && on_true->is_constant() &&
+           on_false->is_constant();
+  }
+  int64_t get_val() const override {
+    return pred->get_val() != 0 ? on_true->get_val() : on_false->get_val();
+  }
+  DynExpr* get_pred() const { return pred.get(); }
+  DynExpr* get_on_true() const { return on_true.get(); }
+  DynExpr* get_on_false() const { return on_false.get(); }
+  DynExpr* substitute(int id, DynExpr* v) override {
+    return new SelectExpr(pred->substitute(id, v), on_true->substitute(id, v),
+                          on_false->substitute(id, v));
+  }
+  std::set<int> get_all_ids() override {
+    auto ids = pred->get_all_ids();
+    ids.merge(on_true->get_all_ids());
+    ids.merge(on_false->get_all_ids());
+    return ids;
+  }
+  std::optional<int64_t> solve(int64_t x) override {
+    StringPrinter printer;
+    print(&printer);
+    LOG(WARNING) << "Cannot solve Select dynamic shape expression for value "
+                 << x << ": " << std::move(printer).ToString();
+    return std::nullopt;
+  }
+  DynExpr* s() override;
+};
+
+DynExpr* operator*(DynExpr& lhs, DynExpr& rhs);
+DynExpr* operator*(int64_t k, DynExpr& rhs);
+DynExpr* operator*(DynExpr& lhs, int64_t k);
+DynExpr* operator/(DynExpr& lhs, DynExpr& rhs);
+DynExpr* operator/(DynExpr& lhs, int64_t d);
+DynExpr* operator+(DynExpr& lhs, DynExpr& rhs);
+DynExpr* operator+(DynExpr& lhs, int64_t d);
+DynExpr* operator-(DynExpr& lhs, DynExpr& rhs);
+DynExpr* operator-(DynExpr& lhs, int64_t d);
+bool operator==(DynExpr& lhs, DynExpr& rhs);
+bool operator==(DynExpr& lhs, int64_t d);
+
+inline DExpr operator*(const DExpr& lhs, const DExpr& rhs) {
+  return DExpr::Adopt(*lhs.get() * *rhs.get());
+}
+inline DExpr operator*(int64_t lhs, const DExpr& rhs) {
+  return DExpr::Adopt(lhs * *rhs.get());
+}
+inline DExpr operator*(const DExpr& lhs, int64_t rhs) {
+  return DExpr::Adopt(*lhs.get() * rhs);
+}
+inline DExpr operator/(const DExpr& lhs, const DExpr& rhs) {
+  return DExpr::Adopt(*lhs.get() / *rhs.get());
+}
+inline DExpr operator/(const DExpr& lhs, int64_t rhs) {
+  return DExpr::Adopt(*lhs.get() / rhs);
+}
+inline DExpr operator+(const DExpr& lhs, const DExpr& rhs) {
+  return DExpr::Adopt(*lhs.get() + *rhs.get());
+}
+inline DExpr operator+(const DExpr& lhs, int64_t rhs) {
+  return DExpr::Adopt(*lhs.get() + rhs);
+}
+inline DExpr operator-(const DExpr& lhs, const DExpr& rhs) {
+  return DExpr::Adopt(*lhs.get() - *rhs.get());
+}
+inline DExpr operator-(const DExpr& lhs, int64_t rhs) {
+  return DExpr::Adopt(*lhs.get() - rhs);
+}
+inline bool operator==(const DExpr& lhs, const DExpr& rhs) {
+  return *lhs.get() == *rhs.get();
+}
+inline bool operator==(const DExpr& lhs, int64_t rhs) {
+  return *lhs.get() == rhs;
+}
+
+inline DExpr DExprFromProto(const xla::ExpressionProto& proto) {
+  switch (proto.node_type_case()) {
+    case ExpressionProto::kConstantValue:
+      return DExpr::Const(proto.constant_value());
+    case ExpressionProto::kVariableId:
+      return DExpr::Var(proto.variable_id());
+    case ExpressionProto::kAddNode: {
+      const auto& add = proto.add_node();
+      return DExprFromProto(add.lhs()) + DExprFromProto(add.rhs());
+    }
+    case ExpressionProto::kSubNode: {
+      const auto& sub = proto.sub_node();
+      return DExprFromProto(sub.lhs()) - DExprFromProto(sub.rhs());
+    }
+    case ExpressionProto::kMulNode: {
+      const auto& mul = proto.mul_node();
+      return DExprFromProto(mul.lhs()) * DExprFromProto(mul.rhs());
+    }
+    case ExpressionProto::kDivNode: {
+      const auto& div = proto.div_node();
+      return DExprFromProto(div.lhs()) / DExprFromProto(div.rhs());
+    }
+    case ExpressionProto::kMaxNode: {
+      const auto& max = proto.max_node();
+      return DExpr::Max(DExprFromProto(max.lhs()),
+                        DExprFromProto(max.rhs()));
+    }
+    case ExpressionProto::kGtNode: {
+      const auto& gt = proto.gt_node();
+      return DExpr::Gt(DExprFromProto(gt.lhs()), DExprFromProto(gt.rhs()));
+    }
+    case ExpressionProto::kSelectNode: {
+      const auto& select = proto.select_node();
+      return DExpr::Select(DExprFromProto(select.pred()),
+                           DExprFromProto(select.on_true()),
+                           DExprFromProto(select.on_false()));
+    }
+    case ExpressionProto::NODE_TYPE_NOT_SET:
+    default:
+      return DExpr::Unknown(kMissingExpressionSentinel);
+  }
+}
+
+inline DynExpr* DynExpr::_(int64_t val) {
+  return new Constant(val);
+}
+inline DynExpr* DynExpr::V(int var_id) { return new Variable(var_id); }
+
+}  // namespace xla
+
+#endif  // XLA_SHAPE_EXPR_H_
diff --git a/third_party/xla/xla/shape_test.cc b/third_party/xla/xla/shape_test.cc
index fb51518b116eb2..e47aaf0d398fb2 100644
--- a/third_party/xla/xla/shape_test.cc
+++ b/third_party/xla/xla/shape_test.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "xla/shape.h"
 
+#include <optional>
+#include <set>
 #include <vector>
 
 #include <gtest/gtest.h>
@@ -22,6 +24,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "xla/hlo/testlib/test.h"
 #include "xla/layout.h"
+#include "xla/printer.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/tsl/platform/test_benchmark.h"
@@ -30,6 +33,12 @@ limitations under the License.
 namespace xla {
 namespace {
 
+std::string DExprToString(const DExpr& expr) {
+  StringPrinter printer;
+  expr->print(&printer);
+  return std::move(printer).ToString();
+}
+
 class ShapeTest : public ::testing::Test {
  protected:
   const Shape opaque_ = ShapeUtil::MakeOpaqueShape();
@@ -46,9 +55,10 @@ class ShapeTest : public ::testing::Test {
   const Shape nested_tuple_ =
       ShapeUtil::MakeTupleShape({tuple_, matrix_, token_});
   const Shape dynamic_matrix_ =
-      ShapeUtil::MakeShape(S32, {5, 2}, {true, false});
+      ShapeUtil::MakeShape(S32, {5, 2}, std::vector<bool>{true, false}, {});
   const Shape unbounded_ =
-      ShapeUtil::MakeShape(F32, {Shape::kUnboundedSize, 784}, {true, false});
+      ShapeUtil::MakeShape(F32, {Shape::kUnboundedSize, 784},
+                           std::vector<bool>{true, false}, {});
 };
 
 // Tests that if the dynamic_dimensions parameter empty in the Shape
@@ -98,8 +108,8 @@ TEST_F(ShapeTest, ShapeToString) {
 }
 
 TEST_F(ShapeTest, DynamicShapeToString) {
-  Shape array_shape =
-      ShapeUtil::MakeShape(F32, {23, 44, 55}, {true, false, true});
+  Shape array_shape = ShapeUtil::MakeShape(
+      F32, {23, 44, 55}, std::vector<bool>{true, false, true}, {});
   EXPECT_EQ("f32[<=23,44,<=55]", array_shape.ToString());
 
   array_shape.set_dynamic_dimension(2, false);
@@ -108,6 +118,177 @@ TEST_F(ShapeTest, DynamicShapeToString) {
   EXPECT_EQ("f32[?,784]", unbounded_.ToString());
 }
 
+TEST_F(ShapeTest, DExprSimplifyScalesMixedDenominators) {
+  DExpr expr = (DExpr::Var(1) / 2) + (DExpr::Var(1) / 3);
+  EXPECT_EQ("((5 * A) / 6)", DExprToString(expr.simplify()));
+}
+
+TEST_F(ShapeTest, DExprSimplifyCombinesEqualFractions) {
+  DExpr expr = (DExpr::Var(1) / 2) + (DExpr::Var(1) / 2);
+  EXPECT_EQ("A", DExprToString(expr.simplify()));
+}
+
+TEST_F(ShapeTest, DExprSolveRejectsZeroDivisor) {
+  DExpr expr = DExpr::Var(1) / DExpr::Const(0);
+  EXPECT_FALSE(expr->solve(7).has_value());
+}
+
+TEST_F(ShapeTest, DExprMaxSimplifiesAndRoundTrips) {
+  DExpr expr = DExpr::Max(DExpr::Var(1), DExpr::Const(4));
+  EXPECT_EQ("max(A, 4)", DExprToString(expr.simplify()));
+
+  DExpr clamped = DExpr::Max(DExpr::Var(1), DExpr::Const(0));
+  EXPECT_EQ("A", DExprToString(clamped.simplify()));
+
+  DExpr positive_affine =
+      DExpr::Max(2 * DExpr::Var(1) - 1, DExpr::Const(0));
+  EXPECT_EQ("((2 * A) - 1)",
+            DExprToString(positive_affine.simplify()));
+
+  DExpr unbounded_below =
+      DExpr::Max(DExpr::Var(1) - DExpr::Var(2), DExpr::Const(0));
+  EXPECT_EQ("max((A + ((-1) * B)), 0)",
+            DExprToString(unbounded_below.simplify()));
+
+  DExpr evaluated = expr.substitute(1, DExpr::Const(7)).simplify();
+  EXPECT_EQ(DExpr::Kind::kConstant, evaluated.kind());
+  EXPECT_EQ(7, evaluated->get_val());
+
+  DExpr divided =
+      DExpr::Max((DExpr::Var(1) + 1) / 2, DExpr::Const(0));
+  DExpr divided_evaluated =
+      divided.substitute(1, DExpr::Const(100)).simplify();
+  EXPECT_EQ(DExpr::Kind::kConstant, divided_evaluated.kind());
+  EXPECT_EQ(50, divided_evaluated->get_val());
+
+  ExpressionProto proto;
+  expr.to_proto(&proto);
+  EXPECT_TRUE(expr == DExprFromProto(proto));
+}
+
+TEST_F(ShapeTest, DExprMaxPartiallySolvesSelectedDynamicBranch) {
+  DExpr expr = DExpr::Max(DExpr::Var(1) - 2, DExpr::Const(0));
+  EXPECT_EQ(100, expr->solve(98));
+
+  DExpr reversed = DExpr::Max(DExpr::Const(0), DExpr::Var(1) - 2);
+  EXPECT_EQ(100, reversed->solve(98));
+
+  // A result equal to the constant bound may come from either branch.
+  EXPECT_FALSE(expr->solve(0).has_value());
+  // A result below the constant bound cannot be produced by Max.
+  EXPECT_FALSE(expr->solve(-1).has_value());
+}
+
+TEST_F(ShapeTest, DExprSelectUsesDynamicPredicate) {
+  DExpr delta = DExpr::Var(1) - 3;
+  DExpr expr = DExpr::Select(DExpr::Gt(delta, DExpr::Const(0)),
+                             DExpr::Const(7), DExpr::Const(11));
+  EXPECT_EQ("select(((A - 3) > 0), 7, 11)",
+            DExprToString(expr.simplify()));
+  EXPECT_EQ(7, expr.substitute(1, DExpr::Const(5))->s()->get_val());
+  EXPECT_EQ(11, expr.substitute(1, DExpr::Const(1))->s()->get_val());
+
+  ExpressionProto proto;
+  expr.to_proto(&proto);
+  EXPECT_TRUE(expr == DExprFromProto(proto));
+}
+
+TEST_F(ShapeTest, DExprUnknownPropagatesThroughGtAndSelect) {
+  DExpr gt = DExpr::Gt(DExpr::Unknown(), DExpr::Const(0)).simplify();
+  EXPECT_TRUE(gt.is_unknown());
+
+  DExpr select =
+      DExpr::Select(DExpr::Unknown(), DExpr::Const(7), DExpr::Const(11))
+          .simplify();
+  EXPECT_TRUE(select.is_unknown());
+}
+
+TEST_F(ShapeTest, DExprConstantDivisionIsNotDynamic) {
+  const DExpr expr = DExpr::Const(7) / DExpr::Const(3);
+
+  EXPECT_FALSE(expr->is_dynamic());
+  EXPECT_EQ(expr->get_val(), 2);
+}
+
+TEST_F(ShapeTest, DExprSubstitutionReplacesEveryVariableOccurrence) {
+  const DExpr a = DExpr::Var(1);
+  const DExpr b = DExpr::Var(2);
+  const DExpr expr = (a + b) * (a - b);
+
+  const DExpr substituted = expr.substitute(1, DExpr::Const(7));
+
+  EXPECT_EQ(substituted, (DExpr::Const(7) + b) * (DExpr::Const(7) - b));
+  const std::set<int> remaining_ids = substituted->get_all_ids();
+  EXPECT_EQ(remaining_ids.size(), 1);
+  EXPECT_EQ(remaining_ids.count(2), 1);
+  EXPECT_EQ(expr->get_all_ids().size(), 2);
+}
+
+TEST_F(ShapeTest, DExprSolveInvertsNestedSingleVariableArithmetic) {
+  const DExpr a = DExpr::Var(1);
+  const DExpr expr = ((3 * a) + 6) / 3;
+
+  const std::optional<int64_t> solved = expr->solve(8);
+
+  ASSERT_TRUE(solved.has_value());
+  EXPECT_EQ(*solved, 6);
+  EXPECT_FALSE((a + DExpr::Var(2))->solve(8).has_value());
+}
+
+TEST_F(ShapeTest, DExprRoundTripPreservesRepeatedVariableStructure) {
+  const DExpr a = DExpr::Var(1);
+  const DExpr b = DExpr::Var(2);
+  const DExpr expr = (a + b) * (a - b);
+  ExpressionProto proto;
+
+  expr.to_proto(&proto);
+  const DExpr decoded = DExprFromProto(proto);
+  ExpressionProto round_tripped_proto;
+  decoded.to_proto(&round_tripped_proto);
+
+  EXPECT_EQ(proto.SerializeAsString(),
+            round_tripped_proto.SerializeAsString());
+}
+
+TEST_F(ShapeTest, DExprFindsSmallestSubexpressionCoveringAllVariables) {
+  const DExpr shared_core = DExpr::Var(1) + DExpr::Var(2);
+  const DExpr expr = (shared_core + 2) * (shared_core + 3);
+
+  EXPECT_EQ(
+      shared_core,
+      expr.find_smallest_subexpression_covering_all_variables());
+}
+
+TEST_F(ShapeTest, DExprFindsCommonCoreAcrossRepeatedBranches) {
+  const DExpr shared_core = DExpr::Var(1) + DExpr::Var(2);
+
+  EXPECT_EQ(shared_core,
+            (shared_core * shared_core)
+                .find_smallest_subexpression_covering_all_variables());
+  EXPECT_EQ(shared_core,
+            (shared_core * (DExpr::Const(3) + shared_core))
+                .find_smallest_subexpression_covering_all_variables());
+}
+
+TEST_F(ShapeTest, DExprUsesParentWhenChildCoresDiffer) {
+  const DExpr complete_expr =
+      (DExpr::Var(1) + DExpr::Var(2)) *
+      (DExpr::Var(1) - DExpr::Var(2));
+
+  EXPECT_EQ(
+      complete_expr,
+      complete_expr.find_smallest_subexpression_covering_all_variables());
+}
+
+TEST_F(ShapeTest, DExprReplacesEveryMatchingSubexpression) {
+  const DExpr shared_core = DExpr::Var(1) + DExpr::Var(2);
+  const DExpr expr = (shared_core + 2) * (shared_core + 3);
+  const DExpr replacement = DExpr::Var(7);
+
+  EXPECT_EQ((replacement + 2) * (replacement + 3),
+            expr.replace_subexpression(shared_core, replacement));
+}
+
 TEST_F(ShapeTest, DeleteDimensions) {
   Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {5, 3, 2, 7, 9},
                                                     {2, 0, 1, 4, 3});
diff --git a/third_party/xla/xla/shape_util.cc b/third_party/xla/xla/shape_util.cc
index 4f92ce19adb1f9..75970dbeceb325 100644
--- a/third_party/xla/xla/shape_util.cc
+++ b/third_party/xla/xla/shape_util.cc
@@ -123,6 +123,7 @@ void PrintBufferShape(Printer* printer, const Shape& shape) {
 // its Layout.
 absl::StatusOr<Shape> MakeShapeWithLayoutInternal(
     PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+    absl::Span<const DExpr> expressions,
     absl::Span<const int64_t> minor_to_major,
     absl::Span<const Tile> tiles, int64_t tail_padding_alignment_in_elements,
     PrimitiveType index_primitive_type, PrimitiveType pointer_primitive_type,
@@ -139,7 +140,8 @@ absl::StatusOr<Shape> MakeShapeWithLayoutInternal(
                            PrimitiveType_Name(element_type));
   }
   TF_ASSIGN_OR_RETURN(Shape shape,
-                      ShapeUtil::MakeValidatedShape(element_type, dimensions));
+                      ShapeUtil::MakeValidatedShape(element_type, dimensions,
+                                                     expressions));
   if (element_size_in_bits ==
       ShapeUtil::ByteSizeOfPrimitiveType(element_type) * 8) {
     // Only set element_size_in_bits if it's different from the default value.
@@ -267,9 +269,20 @@ static std::vector<bool> MakeDynamicDimensions(
   return dynamic_dimensions;
 }
 
+static std::vector<DExpr> MakeExpressions(
+    absl::Span<const int64_t> dimensions) {
+  std::vector<DExpr> expressions;
+  expressions.reserve(dimensions.size());
+  for (int64_t d : dimensions) {
+    expressions.push_back(DExpr::Const(d));
+  }
+  return expressions;
+}
+
 /* static */ Shape ShapeUtil::MakeShape(PrimitiveType element_type,
-                                        absl::Span<const int64_t> dimensions) {
-  return MakeValidatedShape(element_type, dimensions).value();
+                                        absl::Span<const int64_t> dimensions,
+                                        absl::Span<const DExpr> expressions) {
+  return MakeValidatedShape(element_type, dimensions, expressions).value();
 }
 
 /* static */ Shape ShapeUtil::MakeScalarShape(PrimitiveType element_type) {
@@ -278,8 +291,10 @@ static std::vector<bool> MakeDynamicDimensions(
 
 /* static */ Shape ShapeUtil::MakeShape(
     PrimitiveType element_type, absl::Span<const int64_t> dimensions,
-    const std::vector<bool>& dynamic_dimensions) {
-  return MakeValidatedShape(element_type, dimensions, dynamic_dimensions)
+    const std::vector<bool>& dynamic_dimensions,
+    absl::Span<const DExpr> expressions) {
+  return MakeValidatedShape(element_type, dimensions, dynamic_dimensions,
+                            expressions)
       .value();
 }
 
@@ -296,20 +311,37 @@ static std::vector<bool> MakeDynamicDimensions(
 }
 
 /* static */ absl::StatusOr<Shape> ShapeUtil::MakeValidatedShape(
-    PrimitiveType element_type, absl::Span<const int64_t> dimensions) {
-  return MakeValidatedShape(element_type, dimensions,
-                            MakeDynamicDimensions(dimensions));
+    PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+    absl::Span<const DExpr> expressions) {
+  if (expressions.empty() && !dimensions.empty()) {
+    std::vector<DExpr> filled_expressions = MakeExpressions(dimensions);
+    return MakeValidatedShape(element_type, dimensions,
+                              MakeDynamicDimensions(dimensions),
+                              filled_expressions);
+  }
+  return MakeValidatedShape(
+      element_type, dimensions, MakeDynamicDimensions(dimensions), expressions);
 }
 
 /* static */ absl::StatusOr<Shape> ShapeUtil::MakeValidatedShape(
     PrimitiveType element_type, absl::Span<const int64_t> dimensions,
-    const std::vector<bool>& dynamic_dimensions) {
+    const std::vector<bool>& dynamic_dimensions,
+    absl::Span<const DExpr> expressions) {
+  std::vector<DExpr> filled_expressions;
+  if (expressions.empty() && !dimensions.empty()) {
+    filled_expressions = MakeExpressions(dimensions);
+    expressions = filled_expressions;
+  }
   if (dynamic_dimensions.size() != dimensions.size()) {
     return InvalidArgument(
         "dynamic dimensions size %d did not match number of dimensions %d",
         dynamic_dimensions.size(), dimensions.size());
   }
-
+  if (expressions.size() != dimensions.size()) {
+    return InvalidArgument(
+        "expressions size %d did not match number of dimensions %d",
+        expressions.size(), dimensions.size());
+  }
   Shape shape;
   int64_t dense_shape_size = primitive_util::IsArrayType(element_type)
                                  ? primitive_util::ByteWidth(element_type)
@@ -328,6 +360,7 @@ static std::vector<bool> MakeDynamicDimensions(
   for (int i = 0; i < ndims; i++) {
     const int64_t d = dimensions[i];
     const bool is_dynamic = dynamic_dimensions[i];
+    const DExpr& expression = expressions[i];
     if (!Shape::IsValidDimensionSize(d, is_dynamic)) {
       return InvalidArgument("Invalid dimension size %d, is_dynamic=%s", d,
                              is_dynamic ? "true" : "false");
@@ -339,7 +372,7 @@ static std::vector<bool> MakeDynamicDimensions(
       any_overflows |= overflow;
     }
 
-    shape.add_dimensions(d, is_dynamic);
+    shape.add_dimensions(d, is_dynamic, expression);
     minor_to_major->push_back(ndims - 1 - i);
   }
 
@@ -352,11 +385,29 @@ static std::vector<bool> MakeDynamicDimensions(
 
 /* static */ Shape ShapeUtil::MakeShapeWithDenseLayout(
     PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+    absl::Span<const DExpr> expressions,
     absl::Span<const int64_t> minor_to_major, absl::Span<const Tile> tiles,
     int64_t tail_padding_alignment_in_elements, int64_t element_size_in_bits,
     int64_t memory_space, absl::Span<const SplitConfig> split_configs) {
   auto ret = MakeShapeWithLayoutInternal(
-      element_type, dimensions, minor_to_major, tiles,
+      element_type, dimensions, expressions, minor_to_major, tiles,
+      tail_padding_alignment_in_elements,
+      /*index_primitive_type=*/PRIMITIVE_TYPE_INVALID,
+      /*pointer_primitive_type=*/PRIMITIVE_TYPE_INVALID, element_size_in_bits,
+      memory_space, split_configs,
+      /*physical_shape=*/std::nullopt);
+  TF_CHECK_OK(ret.status());
+  return *ret;
+}
+
+/* static */ Shape ShapeUtil::MakeShapeWithDenseLayout(
+    PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+    absl::Span<const int64_t> minor_to_major, absl::Span<const Tile> tiles,
+    int64_t tail_padding_alignment_in_elements, int64_t element_size_in_bits,
+    int64_t memory_space, absl::Span<const SplitConfig> split_configs) {
+  auto ret = MakeShapeWithLayoutInternal(
+      element_type, dimensions, MakeExpressions(dimensions), minor_to_major,
+      tiles,
       tail_padding_alignment_in_elements,
       /*index_primitive_type=*/PRIMITIVE_TYPE_INVALID,
       /*pointer_primitive_type=*/PRIMITIVE_TYPE_INVALID, element_size_in_bits,
@@ -373,7 +424,7 @@ static std::vector<bool> MakeDynamicDimensions(
     int64_t tail_padding_alignment_in_elements, int64_t element_size_in_bits,
     int64_t memory_space, std::optional<Shape> physical_shape) {
   auto ret = MakeShapeWithLayoutInternal(
-      element_type, dimensions, minor_to_major,
+      element_type, dimensions, MakeExpressions(dimensions), minor_to_major,
       /*tiles=*/{}, tail_padding_alignment_in_elements, index_primitive_type,
       pointer_primitive_type, element_size_in_bits, memory_space,
       /*split_configs=*/{}, std::move(physical_shape));
@@ -408,10 +459,11 @@ static std::vector<bool> MakeDynamicDimensions(
 }
 
 /* static */ Shape ShapeUtil::MakeShapeWithDescendingLayout(
-    PrimitiveType element_type, absl::Span<const int64_t> dimensions) {
-  std::vector<int64_t> layout(dimensions.size());
-  std::iota(layout.rbegin(), layout.rend(), static_cast<int64_t>(0));
-  return MakeShapeWithDenseLayout(element_type, dimensions, layout);
+    PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+    absl::Span<const DExpr> expressions) {
+  return MakeShapeWithDenseLayout(
+      element_type, dimensions, expressions,
+      LayoutUtil::MakeDescendingLayout(dimensions.size()).minor_to_major());
 }
 
 /* static */ Shape
@@ -425,7 +477,17 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(
     }
     dims[i] = shape.dimensions(dim);
   }
-  Shape new_shape = MakeShapeWithDescendingLayout(shape.element_type(), dims);
+  std::vector<DExpr> expressions;
+  expressions.reserve(shape.dimensions().size());
+  for (int i = 0; i < shape.dimensions().size(); ++i) {
+    int dim = i;
+    if (shape.has_layout()) {
+      dim = LayoutUtil::Major(shape.layout(), dim);
+    }
+    expressions.push_back(shape.expressions(dim));
+  }
+  Shape new_shape = MakeShapeWithDescendingLayout(shape.element_type(), dims,
+                                                  expressions);
   // Since the physical layout is kept the same, the tiles and element size are
   // the same also.
   if (shape.has_layout()) {
@@ -442,6 +504,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(
       dim = LayoutUtil::Major(shape.layout(), dim);
     }
     new_shape.set_dynamic_dimension(i, shape.is_dynamic_dimension(dim));
+    new_shape.set_expression(i, shape.expressions(dim));
   }
   new_shape.mutable_layout()->set_memory_space(shape.layout().memory_space());
   return new_shape;
@@ -731,7 +794,29 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
         printer->Append("?");
       }
     } else {
+      // Only print constant expression if it is different than the dimension
+      // (i.e. it is wrong!)
+      const DExpr& expr = shape.expressions(i);
+      bool is_wrong = expr && expr->is_constant() &&
+                      expr->get_val() != shape.dimensions(i);
       printer->Append(shape.dimensions(i));
+      if (is_wrong) {
+        DExpr simplified_expr = expr.simplify();
+        xla::StringPrinter expr_printer;
+        simplified_expr->print(&expr_printer);
+        LOG(ERROR) << "Mismatched static shape expression at dim " << i
+                   << ": dim=" << shape.dimensions(i)
+                   << ", expr=" << std::move(expr_printer).ToString();
+        printer->Append("<!");
+        simplified_expr->print(printer);
+        printer->Append("!>");
+      }
+      if (expr && expr->is_dynamic()) {
+        DExpr simplified_expr = expr.simplify();
+        printer->Append("<");
+        simplified_expr->print(printer);
+        printer->Append(">");
+      }
     }
   };
   print_dimension(0);
@@ -753,6 +838,11 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
     return;
   }
   PrintHumanString(printer, shape);
+  if (shape.outer_multiplier() > 0) {
+    printer->Append("(bm=");
+    printer->Append(shape.outer_multiplier());
+    printer->Append(")");
+  }
   if (!shape.IsArray()) return;
   if (!shape.has_layout()) return;
   if (IsScalar(shape)) {
@@ -811,6 +901,10 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
                                             const Shape& rhs) {
   if (!SameRank(lhs, rhs)) return false;
   for (int i = 0; i < lhs.dimensions().size(); ++i) {
+      if (i == 0 && (lhs.outer_multiplier() > 0 || rhs.outer_multiplier() > 0)) {
+        VLOG(3) << "CompareShapes: batch dimension found. Forcely compatible";
+        continue;
+      }
     if (!lhs.is_unbounded_dynamic_dimension(i) &&
         !rhs.is_unbounded_dynamic_dimension(i) &&
         lhs.dimensions(i) != rhs.dimensions(i)) {
@@ -826,7 +920,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
 }
 
 /* static */ bool ShapeUtil::Compatible(const Shape& lhs, const Shape& rhs) {
-  return Shape::Equal().IgnoreDynamicDimension().IgnoreLayout()(lhs, rhs);
+  return Shape::Equal().IgnoreDynamicDimension().IgnoreLayout().IgnoreBatch()(lhs, rhs);
 }
 
 /* static */ bool ShapeUtil::CompatibleIgnoringElementType(const Shape& lhs,
@@ -871,6 +965,11 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) {
   return shape.dimensions(GetDimensionNumber(shape, dimension_number));
 }
 
+/* static */ const DExpr& ShapeUtil::GetExpression(
+    const Shape& shape, int64_t dimension_number) {
+  return shape.expressions(GetDimensionNumber(shape, dimension_number));
+}
+
 /* static */ int64_t ShapeUtil::GetDimensionNumber(const Shape& shape,
                                                    int64_t dimension_number) {
   if (dimension_number < 0) {
@@ -1207,8 +1306,11 @@ ShapeUtil::PackedFactorFor1DInterleavedArray(const Shape& shape) {
   const auto permuted_dims = Permute(shape.dimensions(), permutation);
   const auto permuted_dynamic_dims =
       Permute(shape.dynamic_dimensions(), permutation);
+  const auto permuted_expressions =
+      Permute(shape.expressions(), permutation);
   for (int i = 0; i < permuted_dims.size(); ++i) {
-    new_shape.add_dimensions(permuted_dims[i], permuted_dynamic_dims[i]);
+    new_shape.add_dimensions(permuted_dims[i], permuted_dynamic_dims[i],
+                             permuted_expressions[i]);
   }
 
   // If `shape` has a layout, by contract we choose a new layout such that the
@@ -1763,7 +1865,8 @@ ShapeUtil::DecomposeBitcastToTrt(const Shape& input_shape,
   }
 
   Shape output_shape_with_layout = MakeShapeWithDenseLayout(
-      output_shape.element_type(), output_shape.dimensions(), output_layout);
+      output_shape.element_type(), output_shape.dimensions(),
+      output_shape.expressions(), output_layout);
   CHECK(ReshapeIsBitcast(input_shape, output_shape_with_layout))
       << "reshape is not a bitcast for input_shape: "
       << ShapeUtil::HumanStringWithLayout(input_shape)
@@ -1996,7 +2099,8 @@ struct ParallelState {
   }
 
   // Create the shape of the "work" which has same layout as the original shape.
-  Shape work_shape = ShapeUtil::MakeShape(shape.element_type(), work_dims);
+  Shape work_shape = ShapeUtil::MakeShape(shape.element_type(), work_dims,
+                                          shape.expressions());
   *work_shape.mutable_layout() = shape.layout();
 
   // We target one task (partition) per available thread.
diff --git a/third_party/xla/xla/shape_util.h b/third_party/xla/xla/shape_util.h
index 4e3602671fa153..0948c6c31df2bc 100644
--- a/third_party/xla/xla/shape_util.h
+++ b/third_party/xla/xla/shape_util.h
@@ -330,6 +330,11 @@ class ShapeUtil {
   // GetDimensionNumber(dimension_number).
   static int64_t GetDimension(const Shape& shape, int64_t dimension_number);
 
+  // Extracts the shape's expressions at dimension number
+  // GetDimensionNumber(dimension_number).
+  static const DExpr& GetExpression(const Shape& shape,
+                                    int64_t dimension_number);
+
   // Resolves a dimension number, supporting negative indexing.
   //
   // Negative indexing has similar semantics to Python. For an N-dimensional
@@ -406,7 +411,8 @@ class ShapeUtil {
   // Constructs a new shape with the given element type and sequence of
   // dimensions.
   static Shape MakeShape(PrimitiveType element_type,
-                         absl::Span<const int64_t> dimensions);
+                         absl::Span<const int64_t> dimensions,
+                         absl::Span<const DExpr> expressions = {});
 
   // Make a scalar shape with given primitive type.
   static Shape MakeScalarShape(PrimitiveType element_type);
@@ -419,7 +425,8 @@ class ShapeUtil {
   // the same size.
   static Shape MakeShape(PrimitiveType element_type,
                          absl::Span<const int64_t> dimensions,
-                         const std::vector<bool>& dynamic_dimensions);
+                         const std::vector<bool>& dynamic_dimensions,
+                         absl::Span<const DExpr> expressions = {});
   // Constructs a new buffer shape with the given element type, and sequence of
   // dimensions.
   static Shape MakeBufferShape(PrimitiveType element_type,
@@ -430,10 +437,13 @@ class ShapeUtil {
   // size fits in std::numeric_limits<int64_t>::max(), and dynamic size is not
   // marked static.
   static absl::StatusOr<Shape> MakeValidatedShape(
-      PrimitiveType element_type, absl::Span<const int64_t> dimensions);
+      PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+      absl::Span<const DExpr> expressions = {});
+
   static absl::StatusOr<Shape> MakeValidatedShape(
       PrimitiveType element_type, absl::Span<const int64_t> dimensions,
-      const std::vector<bool>& dynamic_dimensions);
+      const std::vector<bool>& dynamic_dimensions,
+      absl::Span<const DExpr> expressions = {});
 
   // Creates a Shape with element type corresponding to T and the given
   // dimensions
@@ -443,6 +453,17 @@ class ShapeUtil {
                                 dimensions);
   }
 
+  // Constructs a new dense array shape with the given minor_to_major order in
+  // its Layout. Returns a value shape such that shape.has_layout().
+  static Shape MakeShapeWithDenseLayout(
+      PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+      absl::Span<const DExpr> expressions,
+      absl::Span<const int64_t> minor_to_major,
+      absl::Span<const Tile> tiles = {},
+      int64_t tail_padding_alignment_in_elements = 1,
+      int64_t element_size_in_bits = 0, int64_t memory_space = 0,
+      absl::Span<const SplitConfig> split_configs = {});
+
   // Constructs a new dense array shape with the given minor_to_major order in
   // its Layout. Returns a value shape such that shape.has_layout().
   static Shape MakeShapeWithDenseLayout(
@@ -476,7 +497,8 @@ class ShapeUtil {
 
   // Constructs a new shape with major-first layout (i.e. {n, n-1, ..., 0}).
   static Shape MakeShapeWithDescendingLayout(
-      PrimitiveType element_type, absl::Span<const int64_t> dimensions);
+      PrimitiveType element_type, absl::Span<const int64_t> dimensions,
+      absl::Span<const DExpr> expressions = {});
 
   // Returns a new Shape based on the given Shape with low-dimension-major
   // layout (i.e. {n, n-1, ..., 0}, like Fortran), and with the dimensions
diff --git a/third_party/xla/xla/shape_util_test.cc b/third_party/xla/xla/shape_util_test.cc
index 743aa8ecf4d507..5a01598ad7ae1c 100644
--- a/third_party/xla/xla/shape_util_test.cc
+++ b/third_party/xla/xla/shape_util_test.cc
@@ -1049,7 +1049,8 @@ TEST(ShapeUtilTest, InvalidDynamicDimension) {
 TEST(ShapeUtilTest, PermuteDynamicDimensions) {
   Shape shape =
       ShapeUtil::MakeShape(F32, {10, 100, 1000},
-                           /*dynamic_dimensions*/ {false, true, true});
+                           /*dynamic_dimensions*/ {false, true, true},
+                           /*expressions=*/{});
   SCOPED_TRACE(absl::StrCat("shape=", shape.ToString()));
 
   std::vector<int64_t> permutation(3);
@@ -1129,6 +1130,15 @@ TEST(ShapeUtilTest, DeleteDimensions) {
             ShapeUtil::MakeShapeWithDenseLayout(F32, {5, 2}, {1, 0}));
 }
 
+TEST(ShapeUtilTest, MakeShapeWithDenseLayoutPreservesExpressions) {
+  std::vector<DExpr> expressions = {DExpr::Var(7), DExpr::Const(24)};
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(
+      F32, {10, 24}, expressions, {1, 0});
+
+  EXPECT_TRUE(shape.expressions(0) == DExpr::Var(7));
+  EXPECT_TRUE(shape.expressions(1) == DExpr::Const(24));
+}
+
 TEST(ShapeUtilTest, MakeShapeWithDescendingLayoutAndSamePhysicalLayout) {
   Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {128, 24, 4, 48, 48},
                                                     {2, 4, 3, 1, 0});
@@ -1153,6 +1163,24 @@ TEST(ShapeUtilTest,
   EXPECT_EQ(new_shape, expected_shape);
 }
 
+TEST(ShapeUtilTest,
+     MakeShapeWithDescendingLayoutAndSamePhysicalLayoutPreservesExpressions) {
+  std::vector<DExpr> expressions = {DExpr::Var(1), DExpr::Const(24),
+                                    DExpr::Var(2), DExpr::Const(48),
+                                    DExpr::Const(48)};
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(
+      F32, {128, 24, 4, 48, 48}, expressions, {2, 4, 3, 1, 0});
+
+  Shape new_shape =
+      ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(shape);
+
+  EXPECT_TRUE(new_shape.expressions(0) == DExpr::Var(1));
+  EXPECT_TRUE(new_shape.expressions(1) == DExpr::Const(24));
+  EXPECT_TRUE(new_shape.expressions(2) == DExpr::Const(48));
+  EXPECT_TRUE(new_shape.expressions(3) == DExpr::Const(48));
+  EXPECT_TRUE(new_shape.expressions(4) == DExpr::Var(2));
+}
+
 TEST(ShapeUtilTest, DeduceTransposeDimensionsForBitcast) {
   Shape input_shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {5, 3}, {1, 0});
   Shape output_shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {3, 5}, {0, 1});
diff --git a/third_party/xla/xla/stream_executor/tpu/c_api_decl.h b/third_party/xla/xla/stream_executor/tpu/c_api_decl.h
index 0f3ba256573de1..e06db894249632 100644
--- a/third_party/xla/xla/stream_executor/tpu/c_api_decl.h
+++ b/third_party/xla/xla/stream_executor/tpu/c_api_decl.h
@@ -247,6 +247,8 @@ typedef struct XLA_Shape {
   int element_type;
   Int64List dimensions;
   BoolList dynamic_dimensions;
+  Int64List batch_multipliers;
+  Int64List batch_offsets;
   struct XLA_Shape* tuple_shapes;  // owned
   int ntuple_shapes;
   bool has_layout;
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index ca8ba0553bd56a..18abf310ac67e4 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -222,6 +222,8 @@ message DebugOptions {
   // When true, XLA:CPU uses XNNPACK to execute supported operations.
   bool xla_cpu_use_xnnpack = 359;
 
+  string xla_compile_batch_sizes = 399;
+
   // Enabling this will enable optimizations that ignore the possibility of NaN.
   bool xla_enable_fast_math = 335;
 
@@ -1208,7 +1210,7 @@ message DebugOptions {
   // Note: when adding a new flag, please add it to one of the hardware-specific
   // or hardware-agnostic sections at the top of this proto message.
 
-  // Next id: 389
+  // Next id: 400
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.
diff --git a/third_party/xla/xla/xla_data.proto b/third_party/xla/xla/xla_data.proto
index 98462a4f6eb83c..169e82bef57aed 100644
--- a/third_party/xla/xla/xla_data.proto
+++ b/third_party/xla/xla/xla_data.proto
@@ -361,6 +361,8 @@ message ShapeProto {
   // The layout used to back this shape.
   LayoutProto layout = 5;
 
+  repeated ExpressionProto expressions = 7;
+
   // Important: if any field is added, be sure to modify ShapeUtil::Equal(),
   // ShapeUtil::Compatible() and Shape::Hash() appropriately to account for the
   // new field.
@@ -1204,3 +1206,53 @@ message OriginalArrayProto {
 message OriginalValueProto {
   repeated OriginalArrayProto leaves = 1;
 }
+
+message ExpressionProto {
+  oneof node_type {
+    int32 constant_value = 1;  // cons
+    int32 variable_id = 2;     // var
+    AddNode add_node = 3;      // exp + exp
+    SubNode sub_node = 4;      // exp - exp
+    MulNode mul_node = 5;      // exp * exp
+    DivNode div_node = 6;      // exp / exp
+    MaxNode max_node = 7;      // max(exp, exp)
+    GtNode gt_node = 8;        // exp > exp
+    SelectNode select_node = 9;  // select(pred, on_true, on_false)
+  }
+}
+
+message AddNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message SubNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message MulNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message DivNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message MaxNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message GtNode {
+  ExpressionProto lhs = 1;
+  ExpressionProto rhs = 2;
+}
+
+message SelectNode {
+  ExpressionProto pred = 1;
+  ExpressionProto on_true = 2;
+  ExpressionProto on_false = 3;
+}