return;
    }
  }

  // These ops are cheap regardless of their input sizes.
  //
  // TODO(chky): Find a more scalable way to figure out cheap ops.
  if (llvm::isa<mlir::TF::ShapeOp, mlir::TF::StridedSliceOp,
                mlir::TF::ReshapeOp, mlir::TF::ExpandDimsOp>(op)) {
    cost_map_[op] = kDefaultCheapCost;
    return;
  }

  // For other ops, use the sum of input sizes as its cost.

TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b);

// Return an identity op.
TFE_Op* IdentityOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return a shape op fetching the shape of `a`.
TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a);

// Return an allreduce op adding up input tensor `in` from `group_size` workers.
TFE_Op* AllReduceOp(TFE_Context* ctx, TFE_TensorHandle* in, int group_size);

    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());

    TFE_Op* shape_op = ShapeOp(ctx, hgpu);
    TFE_OpSetDevice(shape_op, gpu_device_name.c_str(), status.get());
    ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
    TFE_TensorHandle* retvals[1];
    int num_retvals = 1;
    TFE_Execute(shape_op, &retvals[0], &num_retvals, status.get());

  // CHECK-LABEL:select_v2_with_high_dims_dynamic_shape_both_sides
  // CHECK: %[[SHAPE_0:.*]] = "tfl.shape"(%arg0) : (tensor<8x7x6x5x?x3x2x1xf32>) -> tensor<8xi64>
  // CHECK: %[[SHAPE_1:.*]] = "tfl.shape"(%arg1) : (tensor<?x3x2x1xf32>) -> tensor<4xi64>
  // CHECK: %[[BROADCAST_ARGS_0:.*]] = "tfl.broadcast_args"(%[[SHAPE_0]], %[[SHAPE_1]]) : (tensor<8xi64>, tensor<4xi64>) -> tensor<8xi64>

      mhlo::GetDimensionSizeOp op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
    Value shape_op = rewriter.create<TF::ShapeOp>(op.getLoc(), op.getOperand(),
                                                  rewriter.getBoolAttr(true));
    Value size =
        BuildIntArrayConstOp(builder, rewriter, llvm::SmallVector<int64_t>({1}),

  NativeCodeCall<"GetF32Scalar(&$_builder, " # value # ")">;

def TrueBoolAttr : AttrConstraint<CPred<"$_self.cast<::mlir::BoolAttr>().getValue()">>;

def CreateTFShapeOp : NativeCodeCall<
    "$_builder.create<TF::ShapeOp>($0.getLoc(), $1, $2)">;

def IsI32 : NativeCodeCall<
    "$_builder.getBoolAttr(getElementTypeOrSelf($0.getType()).isInteger(32))">;

def CreateTFCastOpF32 : NativeCodeCall<

    return rewriter.create<TF::ConstOp>(loc, static_shape_attr).getOutput();
  }

  // If the shape is not static, create a new ShapeOp.
  BoolAttr false_attr = rewriter.getBoolAttr(false);
  return rewriter
      .create<TF::ShapeOp>(loc, input,
                           /*use_32bit=*/false_attr)
      .getOutput();
}

mlir::TFL::MirrorPaddingType GetTFLMirrorPaddingFromString(

// CHECK:  [[DIFF_RES:%.*]] = "tf.Greater"([[SIZE_DIFF]], [[ZERO]]) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK:  [[SHAPE_1:%.*]] = "tf.Shape"([[INPUT]]) : (tensor<3x10xf32>) -> tensor<?xi32>

    TypeID::get<TF::ReshapeOp>(),
    TypeID::get<TF::ReverseV2Op>(),
    TypeID::get<TF::RFFTOp>(),
    TypeID::get<TF::RsqrtGradOp>(),
    TypeID::get<TF::ScatterNdOp>(),
    TypeID::get<TF::ShapeOp>(),
    TypeID::get<TF::SinhOp>(),
    TypeID::get<TF::SizeOp>(),
    TypeID::get<TF::SliceOp>(),
    TypeID::get<TF::SoftmaxCrossEntropyWithLogitsOp>(),
    TypeID::get<TF::SoftplusOp>(),

          loc, GetScalarOfType(
                   builder->getIntegerType(use_32bit.getValue() ? 32 : 64),
                   ranked_ty.getDimSize(idx)));
    }
  }

  auto shape = builder->create<TF::ShapeOp>(loc, input, use_32bit);
  return builder->create<TF::StridedSliceOp>(
      loc, mlir::RankedTensorType::get({}, getElementTypeOrSelf(shape)), shape,
      /*begin=*/