// Reshape input image to add a new spatial dimension.
    auto image_type = mlir::cast<ShapedType>(conv_op.getLhs().getType());
    SmallVector<int64_t, 4> image_2d_shape(image_type.getShape().begin(),
                                           image_type.getShape().end());
    image_2d_shape.push_back(1);
    auto image_2d_type =
        RankedTensorType::get(image_2d_shape, image_type.getElementType());
    auto loc = conv_op.getLoc();

    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) {

    SmallVector<int64_t, 4> split_updates_shape;
    split_updates_shape.append(updates_type.getShape().begin(),
                               updates_type.getShape().end());
    split_updates_shape.front() = 1;
    SmallVector<Type, 4> split_updates_type;
    split_updates_type.resize(
        updates_type.getShape().front(),
        tensorflow::GetTypeFromTFTensorShape(split_updates_shape,

      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(),

    if (rhs_shape_type.hasRank()) {
      shape.append(rhs_shape_type.getShape().begin(),
                   rhs_shape_type.getShape().end());
      refined_shape = true;
    }
  } else if (rhs_shape_type.hasRank()) {
    for (auto shape_elts : llvm::enumerate(
             llvm::zip(lhs_shape_type.getShape(), rhs_shape_type.getShape()))) {
      if (ShapedType::isDynamic(std::get<0>(shape_elts.value())) &&

    }
  }

  auto ranked_type = mlir::dyn_cast<RankedTensorType>(result_type);
  if (!ranked_type) return UnrankedTensorType::get(builder->getI1Type());

  return RankedTensorType::get(ranked_type.getShape(), builder->getI1Type());
}

Type InferReductionOpType(Value input, Value reduction_indices,
                          BoolAttr keep_dims) {
  Type input_ty = input.getType();

    const auto [input_height, input_width] =
        GetDimSize(op->getOperand(0).getType().cast<ShapedType>().getShape(),
                   dimension_numbers.getInputSpatialDimensions());
    const auto [output_height, output_width] =
        GetDimSize(op->getResult(0).getType().cast<ShapedType>().getShape(),
                   dimension_numbers.getOutputSpatialDimensions());

def HasEqualShape : Constraint<CPred<
  "$0.getType().cast<ShapedType>().hasRank() && "
  "$1.getType().cast<ShapedType>().hasRank() && "
  "$0.getType().cast<ShapedType>().getShape() == $1.getType().cast<ShapedType>().getShape()">,
  "Checks if the shapes of tensors are same.">;

// Make the 1D value $0 broadcastable with the shape of $1.
def MakeOneDimValueBroadcastable : NativeCodeCall<

      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(