tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output shape_a = ops::Placeholder(s.WithOpName("shape_a"), DT_INT32,
                                    ops::Placeholder::Attrs{});
  Output shape_b = ops::Placeholder(s.WithOpName("shape_b"), DT_INT32,
                                    ops::Placeholder::Attrs{});
  Output shape = ops::Add(s.WithOpName("shape"), shape_a, shape_b);

      shape_before.take_back(shape_c.size()) != shape_c) {
    return failure();
  }
  ArrayRef<int64_t> shape_y2 =
      shape_before.drop_front(shape_b.size()).drop_back(shape_c.size());

  // No need to check %dot; dot_general verifier ensures correct shapes.
  // Extract Z from %dot.
  ArrayRef<int64_t> shape_z =
      dot.getType().getShape().drop_front(shape_b.size() + shape_y2.size());

  const TF_Shape* shape_b = TF_TensorSpecShape(tensor_spec_b);

  // Input "b" is a scalar, float32 tensor
  EXPECT_EQ("b", std::string(TF_SignatureDefParamName(param_b)));
  EXPECT_EQ(TF_FLOAT, TF_TensorSpecDataType(tensor_spec_b));
  EXPECT_EQ(0, TF_ShapeDims(shape_b));

  EXPECT_EQ(TF_SignatureDefParamListSize(returns), 1);

  const TF_SignatureDefParam* param_out =

  // CHECK-DAG: %[[SHAPE1:.*]] = "tf.Const"() <{value = dense<[1, 32, 32, 16]> : tensor<4xi64>}>
  // CHECK: %[[SHAPE2:.*]] = "tf.Shape"(%arg2) : (tensor<*xf32>) -> tensor<?xi64>
  %0:3 = "tf.ShapeN"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<1x32x32x16xf32>, tensor<*xf32>) -> (tensor<0xi64>, tensor<4xi64>, tensor<?xi64>)

  // CHECK: return %[[SHAPE0]], %[[SHAPE1]], %[[SHAPE2]]

// CHECK-DAG: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]] : tensor<?x3x2x2xf32> -> tensor<4xindex>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST_0]], %[[SHAPE_OF]], dims = [3] : (tensor<2xf32>, tensor<4xindex>) -> tensor<?x3x2x2xf32>

  // CHECK-DAG:  %[[CSTR_LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:  %[[CSTR_RHS_SHAPE:.+]] = shape.shape_of %arg1
  // CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[CSTR_LHS_SHAPE]], %[[CSTR_RHS_SHAPE]]
  // CHECK-NEXT: shape.assuming %[[WITNESS:.+]]
  // CHECK-DAG:    %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
  // CHECK-DAG:    %[[RHS_SHAPE:.+]] = shape.shape_of %arg1

// RUN: tf-opt -split-input-file -verify-diagnostics --tf-shape-inference %s | FileCheck %s

module attributes {tf.versions = {producer = 888 : i32}} {
// CHECK-LABEL: testConv2dShapeValidPadding
func.func @testConv2dShapeValidPadding(%arg0: tensor<1x112x80x128xf32>, %arg1: tensor<128x3x3x128xf32>, %arg2: tensor<128xf32>) -> tensor<1x?x?x128xf32> {

  ParallelTensor(const ParallelDevice& device,
                 std::vector<TensorHandlePtr> tensors,
                 absl::Span<const int64_t> shape, const TF_DataType dtype)
      : device_(device),
        tensors_(std::move(tensors)),
        shape_(std::vector<int64_t>(shape.begin(), shape.end())),
        dtype_(dtype) {}
  ParallelTensor(const ParallelDevice& device,

    %mean: tensor<?xf32>, %variance: tensor<?xf32>)
    -> tensor<?x?x?x?xf32> {
  // CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e-03> : tensor<f32>
  // CHECK-DAG: %[[VAR_SHAPE:.+]] = shape.shape_of %[[VARIANCE]] : tensor<?xf32> -> tensor<1xindex>

  for (auto operand_and_idx : llvm::enumerate(cluster_func.getOperands())) {
    // Skip adding shape op for operands that have static shapes.
    tensorflow::PartialTensorShape shape(
        metadata.args(operand_and_idx.index()).shape());
    if (shape.IsFullyDefined()) continue;

    auto shape_op = builder->create<TF::ShapeOp>(
        cluster_func.getLoc(),