| 34 | } |
| 35 | |
| 36 | xla::StatusOr<PrimitiveType> DtypeToPrimitiveType(const py::dtype& np_type) { |
| 37 | static auto* types = |
| 38 | new absl::flat_hash_map<std::pair<char, int>, PrimitiveType>({ |
| 39 | {{'b', 1}, PRED}, |
| 40 | {{'i', 1}, S8}, |
| 41 | {{'i', 2}, S16}, |
| 42 | {{'i', 4}, S32}, |
| 43 | {{'i', 8}, S64}, |
| 44 | {{'u', 1}, U8}, |
| 45 | {{'u', 2}, U16}, |
| 46 | {{'u', 4}, U32}, |
| 47 | {{'u', 8}, U64}, |
| 48 | {{'V', 2}, BF16}, // array protocol code for raw data (void*) |
| 49 | {{'f', 2}, F16}, |
| 50 | {{'f', 4}, F32}, |
| 51 | {{'f', 8}, F64}, |
| 52 | {{'c', 8}, C64}, |
| 53 | {{'c', 16}, C128}, |
| 54 | }); |
| 55 | auto it = types->find({np_type.kind(), np_type.itemsize()}); |
| 56 | if (it == types->end()) { |
| 57 | return InvalidArgument("Unknown NumPy type %c size %d", np_type.kind(), |
| 58 | np_type.itemsize()); |
| 59 | } |
| 60 | return it->second; |
| 61 | } |
| 62 | |
| 63 | xla::StatusOr<py::dtype> PrimitiveTypeToDtype(PrimitiveType type) { |
| 64 | switch (type) { |
no test coverage detected