if (mlir::isa<mlir::TF::StringType>(key_type.getElementType()))
    cost *= kLookupTableFindStringKeyCostScale;

  return cost;
}

// The cost function for tf.GatherV2.
int64_t InferGatherV2Cost(const CostContext& context, mlir::TF::GatherV2Op op) {
  return InferTensorSize(
      context, mlir::cast<mlir::TensorType>(op.getOutput().getType()));
}

// The cost function for tf.SparseSegmentSumOp.