| 61 | } |
| 62 | |
| 63 | xla::StatusOr<py::dtype> PrimitiveTypeToDtype(PrimitiveType type) { |
| 64 | switch (type) { |
| 65 | case PRED: |
| 66 | return py::dtype::of<bool>(); |
| 67 | case S8: |
| 68 | return py::dtype::of<int8>(); |
| 69 | case S16: |
| 70 | return py::dtype::of<int16>(); |
| 71 | case S32: |
| 72 | return py::dtype::of<int32>(); |
| 73 | case S64: |
| 74 | return py::dtype::of<int64>(); |
| 75 | case U8: |
| 76 | return py::dtype::of<uint8>(); |
| 77 | case U16: |
| 78 | return py::dtype::of<uint16>(); |
| 79 | case U32: |
| 80 | return py::dtype::of<uint32>(); |
| 81 | case U64: |
| 82 | return py::dtype::of<uint64>(); |
| 83 | case BF16: { |
| 84 | TF_ASSIGN_OR_RETURN(py::object bfloat16, Bfloat16Dtype()); |
| 85 | return py::dtype::from_args(bfloat16); |
| 86 | } |
| 87 | case F16: |
| 88 | return py::dtype("e"); // PEP 3118 code for "float16 |
| 89 | case F32: |
| 90 | return py::dtype::of<float>(); |
| 91 | case F64: |
| 92 | return py::dtype::of<double>(); |
| 93 | case C64: |
| 94 | return py::dtype::of<std::complex<float>>(); |
| 95 | case C128: |
| 96 | return py::dtype::of<std::complex<double>>(); |
| 97 | default: |
| 98 | return Unimplemented("Unimplemented primitive type %s", |
| 99 | PrimitiveType_Name(type)); |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | // Returns a numpy-style format descriptor string for `type`. |
| 104 | StatusOr<std::string> FormatDescriptorForPrimitiveType(PrimitiveType type) { |
no test coverage detected