dyn_cast_or_null<TFL::ExpOp>(sum_op.getInput().getDefiningOp());
    if (!exp_op || !exp_op->hasOneUse()) {
      return failure();
    }

    auto parent_sub_op =
        dyn_cast_or_null<TFL::SubOp>(sub_op.getLhs().getDefiningOp());
    if (!parent_sub_op || parent_sub_op != dyn_cast_or_null<TFL::SubOp>(
                                               exp_op.getX().getDefiningOp())) {
      return failure();
    }

        add_op != nullptr &&
        !isa<stablehlo::ConstantOp>(add_op->getOperand(1).getDefiningOp())) {
      LLVM_DEBUG(llvm::dbgs() << "Expected a `stablehlo.constant` as the "
                              << "rhs of `stablehlo.add`.\n");
    }

    // Make sure the filter is a constant or a constant transpose.
    Operation* filter_op = filter.getDefiningOp();
    const bool is_constant = isa_and_nonnull<stablehlo::ConstantOp>(filter_op);

    Operation* arg_op = arg.getDefiningOp();
    return arg_op && arg_op->getName() == first_arg_op->getName();
  });
  if (!args_same_op) return failure();

  // Collect unary operations operands.
  auto unary_operands = llvm::map_range(op.getValues(), [](Value arg) -> Value {
    return arg.getDefiningOp()->getOperand(0);
  });

  while (
      llvm::isa_and_nonnull<IdentityOp, IdentityNOp>(value.getDefiningOp())) {
    value = value.getDefiningOp()->getOperand(
        mlir::cast<OpResult>(value).getResultNumber());
  }

  Operation* op = value.getDefiningOp();
  if (!llvm::isa_and_nonnull<BatchDatasetV2Op, MapDatasetOp, RepeatDatasetOp,

  auto slice_op =
      dyn_cast_or_null<StridedSliceOp>(getValues()[0].getDefiningOp());
  if (!slice_op) return {};

  // Input to the slice op is defined by shape operation.
  auto shape_op =
      dyn_cast_or_null<ShapeOp>(slice_op.getInput().getDefiningOp());
  if (!shape_op) return {};

  // Input tensor, which shape is reconstructed by the pack operation.

  auto* defining_op = operand.getDefiningOp();
  if (!llvm::dyn_cast_or_null<TF::BroadcastToOp>(defining_op) &&
      !llvm::dyn_cast_or_null<TFL::BroadcastToOp>(defining_op)) {
    return nullptr;
  }

  Value broadcast_shape = defining_op->getOperand(
      1);  // Broadcasted shape operand of BroadcastTo op.
  Operation* parent_of_defining_op = broadcast_shape.getDefiningOp();

  // None type is allowed to represent unspecified operands.
  if (mlir::isa<NoneType>(value.getType())) return true;

  auto type = mlir::dyn_cast<TensorType>(value.getType());
  if (!type) {
    if (auto op = value.getDefiningOp()) {
      error_handler.emitError()
          << '\'' << op << "' should produce value of tensor type instead of "
          << value.getType();
      return false;
    }

  // Returns the first producer op whose type is not in Tys.
  template <typename... Tys>
  Value recursivelyWalkUp(Value op) const {
    while (llvm::isa_and_nonnull<Tys...>(op.getDefiningOp())) {
      Operation* producer = op.getDefiningOp();
      op = producer->getOperand(/*idx=*/0);
    }

    return op;
  }
};

class ConvertMaxPoolOp : public OpConversionPattern<mhlo::ReduceWindowOp> {

                                 int dst_input) {
  // Get the NextIteration.Source operation from the token operand of the sink.
  mlir::Operation* source = sink->getOperand(0).getDefiningOp();

  // Adds the "source" to the operands of the dst by creating a new dst
  // operation.
  mlir::OperationState state(dst->getLoc(), dst->getName());
  auto num_operands = dst->getNumOperands();