rhs_shape.reserve(dnums.lhs_rhs_out.size() + 2);
  for (auto i : dnums.lhs_rhs_out) {
    const int64_t b1 = lhs_type.getShape()[std::get<0>(i)];
    lhs_shape.push_back(b1);
    const int64_t b2 = rhs_type.getShape()[std::get<1>(i)];
    rhs_shape.push_back(b2);

    lhs_reshape_segids.at(std::get<0>(i)) = num_lhs_reshape_segids++;

LogicalResult HandlePad(TF::PadOp op, int32_t kernel_size, int32_t block_size) {
  auto ranked_type = mlir::dyn_cast<RankedTensorType>(op.getInput().getType());
  if (!ranked_type) return failure();
  auto pad_input_shape = ranked_type.getShape();
  Location loc = op.getLoc();
  OpBuilder builder(op);
  builder.setInsertionPoint(op);
  auto padding_type = RankedTensorType::get({4, 2}, builder.getIntegerType(32));

  // Calculate paddings.

    const int num_dimensions = input_tensor.getShapedType().getRank();
    assert(perm_tensor.getType().getNumElements() == num_dimensions);

    ArrayRef<int64_t> input_shape = input_tensor.getShapedType().getShape();
    auto output_type = mlir::cast<ShapedType>(op.getOutput().getType());

    SmallVector<int32_t, 4> perm;
    SmallVector<int64_t, 4> output_shape;
    for (int i = 0; i < num_dimensions; ++i) {

      return failure();
    }

    auto current_out_type =
        llvm::dyn_cast<RankedTensorType>(op->getResult(0).getType());
    auto new_out_type = RankedTensorType::get(
        PermuteShape(current_out_type.getShape(), inverse_perm_arr),
        current_out_type.getElementType());

    // Create new ewise op to appear before the tranpose.
    auto *new_ewise_op =
        rewriter.create(op->getLoc(), op->getName().getIdentifier(),

      auto scatter_dims_to_operand_dims =
          scatter_dimension_numbers.getScatterDimsToOperandDims();

      if (IsIotaAttr(inserted_window_dims, indices_type.getShape().back()) &&
          IsIotaAttr(scatter_dims_to_operand_dims,
                     indices_type.getShape().back())) {
        rewriter.replaceOpWithNewOp<TfOp>(scatter_op,
                                          scatter_op.getResult(0).getType(),

  auto input_type = mlir::cast<mlir::TensorType>(src_input.getType());

  if (input_type.hasRank()) {
    if (input_type.getShape()[split_dimension] == mlir::ShapedType::kDynamic) {
      output_type = input_type;
    } else {
      auto shape = llvm::to_vector<4>(input_type.getShape());
      if (shape[split_dimension] % num_split != 0) {
        return mlir::emitError(
            location,
            llvm::formatv(

        type, builder.getFloatAttr(element_ty, unique_index));

  if (auto qtype = mlir::dyn_cast<QuantizedType>(element_ty)) {
    mlir::RankedTensorType new_type = tensorflow::GetTypeFromTFTensorShape(
        type.getShape(), qtype.getStorageType());
    return DenseElementsAttr::get(
        new_type, builder.getIntegerAttr(qtype.getStorageType(), unique_index));
  }
  llvm_unreachable("unhandled element type");
}

    DenseFPElementsAttr attr;
    if (!matchPattern(weight_op->getResult(0), m_Constant(&attr))) {
      return failure();
    }

    // Get new shape.
    llvm::ArrayRef<int64_t> cur_shape = attr.getType().getShape();
    int cur_rank = cur_shape.size();
    if (cur_rank != 4 || cur_shape[2] == 1) return failure();
    TensorType new_shape = RankedTensorType::get(
        {cur_shape[0], cur_shape[1], 1, cur_shape[2] * cur_shape[3]},

    Value bcast_op_result = (*bcast_op)->getResult(0);
    auto bcast_op_result_type =
        mlir::cast<RankedTensorType>(bcast_op_result.getType());
    const ArrayRef<int64_t> bcast_shape = bcast_op_result_type.getShape();
    const TensorType new_bcast_op_result_type = bcast_op_result_type.cloneWith(
        bcast_shape, accumulation_quantized_element_type);
    bcast_op_result.setType(new_bcast_op_result_type);
  }

    return PartialTensorShape();
  }

  if (auto tensor_type = mlir::dyn_cast<mlir::RankedTensorType>(type)) {
    TensorShapeProto tensor_shape_proto;
    ConvertToTensorShapeProto(tensor_type.getShape(), &tensor_shape_proto);
    return PartialTensorShape(tensor_shape_proto);
  }

  // If type is not a RankedTensor or UnrankedTensor, it must be a scalar.
  // Empty TensorShape indicates a scalar.