auto value_to_replace = avg_pool.getResult();

  // If it's not keep dim, we need to insert a reshape after the average
  // pool.
  if (!mean_op.getKeepDims()) {
    // Insert the reshape.
    SmallVector<int64_t, 2> new_shape({batch, channel});
    auto reshape_op =
        InsertReshapeOp(mean_op.getLoc(), avg_pool.getResult(),

    if (!log_op || !log_op->hasOneUse()) {
      return failure();
    }
    auto sum_op = dyn_cast_or_null<TFL::SumOp>(log_op.getX().getDefiningOp());
    if (!sum_op || !sum_op.getKeepDims() ||
        !isSupportedAxis(
            sum_op.getAxes(),
            mlir::cast<ShapedType>(sum_op.getOperand(0).getType()).getRank())) {
      return failure();
    }
    if (!sum_op->hasOneUse()) {

  if (!result_ty) return {};

  // Bypass this op if the result has the same shape and type. This can happen
  // if the input tensor has size 0 or size 1.
  if (!getKeepDims() && input_ty == result_ty) {
    return getInput();
  }
  return {};
}

//===----------------------------------------------------------------------===//
// StridedSliceOp

    // reshape. Various code generation techniques benefit from the knowledge
    // that this is a restricted form of shape manipulation that is just adding
    // unit dims.
    if (op.getKeepDims()) {
      for (const auto &dim_is_reduced :
           llvm::enumerate(reduced_dimensions_bitmap)) {
        if (dim_is_reduced.value()) {
          auto index_attr = GetI32ElementsAttr(