StringRef config = recv_op ? recv_op.getConfig() : send_op.getConfig();

  // Create XlaRecvTPUEmbeddingDeduplicationData op.
  OpBuilder builder(region);
  auto output_ty =
      RankedTensorType::get({}, VariantType::get(region->getContext()));
  auto dedup_op = builder.create<XlaRecvTPUEmbeddingDeduplicationDataOp>(
      loc, output_ty, config);

  // Rewrite RecvTPUEmbeddingActivations op to the corresponding internal op.

      return builder.getIntegerType(64, /*isSigned=*/false);
    case tflite::TensorType_RESOURCE:
      return mlir::TF::ResourceType::get(builder.getContext());
    case tflite::TensorType_VARIANT:
      return mlir::TF::VariantType::get(builder.getContext());
  }
}

tensorflow::DataType TflTypeToTfType(tflite::TensorType type) {
  switch (type) {
    case tflite::TensorType_BOOL:
      return tensorflow::DT_BOOL;

  // require "raising" later.
  for (const Type type : inst->getResultTypes()) {
    if (const TensorType tensor_type = mlir::dyn_cast<TensorType>(type)) {
      if (mlir::isa<VariantType>(tensor_type.getElementType())) {
        return false;
      }
    }
  }

  // Operations that execute function calls shouldn't be constant folded.