}

PermutationAndShape GetPermutationAndTransposedShape(
    llvm::ArrayRef<int64_t> permutation_array, ShapedType input_type,
    ConversionPatternRewriter& rewriter) {
  assert(permutation_array.size() == input_type.getRank());
  llvm::SmallVector<int64_t> transposed_shape(permutation_array.size());
  for (int64_t i = 0; i < permutation_array.size(); ++i) {
    transposed_shape[i] = input_type.getDimSize(permutation_array[i]);
  }

  if (!lhs_type.hasRank() || !rhs_type.hasRank()) {
    return rewriter.notifyMatchFailure(op, "unsupported unranked input type");
  }
  if (lhs_type.getRank() < 1 || 2 < lhs_type.getRank() ||
      rhs_type.getRank() < 1 || 2 < rhs_type.getRank()) {
    return rewriter.notifyMatchFailure(
        op,
        "unsupported dot operation type; operands must be vectors or "
        "matrices");
  }

  if (!dims_type) return success();
  if (dims_type.getRank() > 1)
    return emitError(loc, "dimensions can only be 0D or 1D tensor");

  auto input_type = mlir::dyn_cast<RankedTensorType>(input.getType());
  if (!input_type) return success();
  int64_t rank = input_type.getRank();

  DenseIntElementsAttr dims_attr;
  if (!matchPattern(dims, m_Constant(&dims_attr))) return success();

  const int64_t input_rank =
      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(0).getType())
          .getRank();
  const int64_t filter_rank =
      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(1).getType())
          .getRank();
  // The following conditions are such requirements:
  //   - rank(lhs) is 1 or 2
  //   - rank(rhs) = 2
  //   - size(lhs_contracting_dimensions) = 1

          permutation_and_shape.shape.getRank() - inserted_window_dims.size();
      int64_t num_updates = indices_type.getDimSize(0);
      // For TF::TensorScatterUpdateOp, `indices` must have at least 2 axes:
      // `(num_updates, index_depth)`. Reshape indices and updates if necessary.
      if (std::is_same<TfOp, TF::TensorScatterUpdateOp>::value &&
          indices_type.getRank() == 1 && updates_type.getRank() == 1 &&

        llvm::dyn_cast<RankedTensorType>(tpose_arg1->getResultTypes()[0]);
    auto tpose_arg2_type =
        llvm::dyn_cast<RankedTensorType>(tpose_arg2->getResultTypes()[0]);
    if (tpose_arg1_type.getRank() != tpose_arg2_type.getRank()) {
      return failure();
    }

    if (llvm::isa<BlockArgument>(tpose_arg1.getPerm()) ||
        llvm::isa<BlockArgument>(tpose_arg2.getPerm())) {
      return failure();
    }

    // Compute mean
    int64_t input_last_dim = input_type.getRank() - 1;
    auto dims_type = RankedTensorType::get(/*shape=*/{input_last_dim},
                                           rewriter.getIntegerType(32));
    ::mlir::SmallVector<int32_t> reduce_dim_axes;
    for (int i = 0; i < input_type.getRank(); ++i) {
      if (i != feature_index) {
        reduce_dim_axes.push_back(i);
      }
    }

    // before broadcasting.
    if (operand_type.getRank() < output_type.getRank()) {
      input = InsertExpandDimsOp(op, rewriter, input, output_type.getRank());
    }

    SmallVector<int32_t> broadcast_shape =
        CastI64ArrayToI32(output_type.getShape()).value();
    TensorType broadcast_shape_type =
        output_type.cloneWith({output_type.getRank()}, rewriter.getI32Type());
    auto broadcast_shape_attr =

  if (!image_type || !image_type.getElementType().isF32() ||
      image_type.getRank() != 4) {
    return func_.emitWarning() << "Image should be a 4D float tensor";
  }

  auto flow_type = mlir::dyn_cast_or_null<RankedTensorType>(
      func_.getFunctionType().getInput(1));
  if (!flow_type || !flow_type.getElementType().isF32() ||
      flow_type.getRank() != 4) {
    return func_.emitWarning() << "Flow should be a 4D float tensor";

  if (!output_ty) return success();

  int64_t expected_output_rank = std::max(x_ty.getRank(), y_ty.getRank());
  if (output_ty.getRank() != expected_output_rank)
    return op.emitOpError()
           << "found invalid output rank, expected " << expected_output_rank
           << " but got " << output_ty.getRank();

  // Check output batch dim with potential broadcasting.