if (parser.parseKeyword(&keyword)) return Type();

  if (keyword == "tf_tensor") return TFTensorType::get(getContext());
  if (keyword == "tf_allocator") return TFAllocatorType::get(getContext());

  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
  return Type();
}

/// Print a type registered to this dialect.
void FallbackDialect::printType(Type type, DialectAsmPrinter &os) const {

                                    BlockArgument& block_arg) {
  // check that inputs length is same as num_replicas.
  if (inputs.size() != replicate.getN()) {
    return replicate.emitError()
           << "Expected numper of inputs (" << inputs.size()
           << ") to append to replicate to be num_replicas ("
           << replicate.getN() << ")";
  }

    if (!sym) continue;
    auto callee = symtab.lookup<func::FuncOp>(sym.getRootReference());
    if (!callee) {
      // This is not expected to happen in practice.
      return op->emitError()
             << "Cannot find function " << sym.getRootReference();
    }
    callees.push_back(callee);
  }
  return success();
}

bool HasSingleBlock(func::FuncOp func) {

    if (op.intermediates.size() != 5) {
      auto err = errors::InvalidArgument(
          "operator has intermediate tensors but the number of them is not "
          "five.");
      return emitError(loc, err.ToString()), err;
    }
    // Create intermediate value

    const llvm::SmallVector<llvm::StringRef, 5> kIntermediateNames = {
        "input_to_input_intermediate", "input_to_forget_intermediate",

// 2. The wrapped launch op should be placed on CPU.
LogicalResult CheckOpIsValid(Operation* op) {
  auto launch_op = llvm::dyn_cast<tf_device::LaunchOp>(op->getParentOp());
  if (!launch_op) {
    op->emitError() << "TPUCopyWithDynamicShapeOp is not in a launch";
  }
  std::string device_str = launch_op.getDeviceAttr().getValue().str();
  std::string cpu0_device;
  if (failed(tensorflow::GetNonReplicatedCPU0(op, &cpu0_device)))

    if (!replication_info_attr)
      return metadata_op.emitError() << kBadReplicateInfoAttrMsg;

    auto replication_info_attr_str =
        mlir::dyn_cast<StringAttr>(replication_info_attr);
    if (!replication_info_attr_str ||
        replication_info_attr_str.getValue().empty())
      return metadata_op.emitError() << kBadReplicateInfoAttrMsg;

    // Remove `name` attribute.

  target.addLegalDialect<arith::ArithDialect>();

  if (failed(applyPartialConversion(getOperation(), target,
                                    std::move(patterns)))) {
    getOperation().emitError("Composite lowering pass failed.");
    signalPassFailure();
  }
}

}  // namespace

// Creates an instance of the pass.
std::unique_ptr<OperationPass<ModuleOp>> CreateCompositeLoweringPass() {

    if (std::find_if(
            attr_to_op_map.begin(), attr_to_op_map.end(), [&](auto attr_op) {
              return std::get<0>(attr_op).getName() == attribute.getName();
            }) == attr_to_op_map.end()) {
      emitError(UnknownLoc::get(&context),
                "Could not find attribute: " + attribute.getName().str());
      return failure();
    }

    Operation* owner_op;
    for (const auto& [attr, val] : attr_to_op_map) {

        new_arg = block.addArgument(arg_type, loc);
        quantize_op.setOperand(new_arg);
      } else {
        input_type.print(llvm::errs() << "Requested input type ");
        quantize_op.emitError(" Couldn't be modified to the requested type.");
        return failure();
      }
      new_input_types[i] = arg_type;
      arg.dropAllUses();
      if (quantize_op.use_empty()) {
        quantize_op.erase();

  ReplicateOp::ensureTerminator(body, parser.getBuilder(), result.location);

  if (!llvm::hasSingleElement(body))
    return parser.emitError(loc) << "expects a single block region";

  Operation& terminator = body.front().back();
  if (!isa<ReturnOp>(terminator))
    return parser.emitError(loc) << "expects a tf_device.return terminator";

  // Get the results type from the terminator type inside the replicate,