// op. The MapDataset op always has N+1 inputs.
  // TODO(jpienaar): Avoid this lookup.
  auto module = op->getParentOfType<ModuleOp>();
  auto f = module.lookupSymbol<func::FuncOp>(op.getF());
  // Skip if function is not found or more than one caller.
  if (!f || !llvm::hasSingleElement(GetCallers(f))) return false;
  return InferShapeForDatasetOpCommon(op, f, max_iterations);
}

  let extraClassDeclaration = [{
    // Returns the callee of this operation.
    func::FuncOp func() {
      return SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(*this, getF());
    }

    // Erases variable arguments from `batch_func_op`. `erase_indices` contains
    // the indices of the arguments to erase.
    void eraseArguments(const BitVector& erase_indices);