tensorflow::DataType XrtComputationClient::XlaTypeToDataType()

in Sources/x10/xla_client/xrt_computation_client.cc [2043:2081]


tensorflow::DataType XrtComputationClient::XlaTypeToDataType(
    PrimitiveType dtype) {
  switch (dtype) {
    case PrimitiveType::PRED:
      return tensorflow::DT_BOOL;
    case PrimitiveType::S8:
      return tensorflow::DT_INT8;
    case PrimitiveType::U8:
      return tensorflow::DT_UINT8;
    case PrimitiveType::S16:
      return tensorflow::DT_INT16;
    case PrimitiveType::U16:
      return tensorflow::DT_UINT16;
    case PrimitiveType::S32:
      return tensorflow::DT_INT32;
    case PrimitiveType::U32:
      return tensorflow::DT_UINT32;
    case PrimitiveType::S64:
      return tensorflow::DT_INT64;
    case PrimitiveType::U64:
      return tensorflow::DT_UINT64;
    case PrimitiveType::F32:
      return tensorflow::DT_FLOAT;
    case PrimitiveType::F64:
      return tensorflow::DT_DOUBLE;
    case PrimitiveType::BF16:
      return tensorflow::DT_BFLOAT16;
    case PrimitiveType::F16:
      return tensorflow::DT_HALF;
    case PrimitiveType::C64:
      return tensorflow::DT_COMPLEX64;
    case PrimitiveType::C128:
      return tensorflow::DT_COMPLEX128;
    default:
      break;
  }
  XLA_ERROR() << "Unable to convert XLA type " << dtype
              << " to tensorflow DataType";
}