| 80 | // transpose_output feature. |
| 81 | template <typename T> |
| 82 | inline XlaOp Dequantize(XlaOp input, const QuantizedRange& range, |
| 83 | absl::string_view mode_string = "MIN_COMBINED", |
| 84 | bool transpose_output = false) { |
| 85 | XlaBuilder* const builder = input.builder(); |
| 86 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 87 | float half_range = |
| 88 | !std::is_signed<T>::value |
| 89 | ? 0.0f |
| 90 | : (static_cast<float>(std::numeric_limits<T>::max()) - |
| 91 | std::numeric_limits<T>::min() + 1) / |
| 92 | 2.0f; |
| 93 | const int64 unpack_size = sizeof(uint32) / sizeof(T); |
| 94 | TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(input)); |
| 95 | |
| 96 | auto element_type = shape.element_type(); |
| 97 | if (element_type != U32) { |
| 98 | return InvalidArgument( |
| 99 | "Only U32 is supported for input type of xla::Dequantize Op."); |
| 100 | } |
| 101 | |
| 102 | // Broadcast the input to [unpack_size, d0, ..., dn] if input size is |
| 103 | // [d0, ..., dn]. |
| 104 | auto broadcast_input = Broadcast(input, {unpack_size}); |
| 105 | |
| 106 | XlaOp iota_r1 = Iota(builder, U32, unpack_size); |
| 107 | // Highest significant bytes needs to shift more bytes than lower |
| 108 | // significant bytes. |
| 109 | XlaOp shift_bytes = |
| 110 | xla::ConstantR0<uint32>(builder, unpack_size - 1) - iota_r1; |
| 111 | |
| 112 | const int bytes_of_type = sizeof(T) / sizeof(uint8); |
| 113 | std::vector<uint32> shift_vec(unpack_size, kBitsOfByte * bytes_of_type); |
| 114 | XlaOp shift_bits = |
| 115 | shift_bytes * xla::ConstantR1<uint32>(builder, shift_vec); |
| 116 | |
| 117 | // Make bit_mask for different data type T. |
| 118 | uint32 bit_mask = 0x00000000; |
| 119 | for (int i = 0; i < bytes_of_type; i++) { |
| 120 | bit_mask <<= kBitsOfByte; |
| 121 | bit_mask |= 0x000000ff; |
| 122 | } |
| 123 | |
| 124 | std::vector<int64> shift_transpose_dimensions(shape.dimensions_size()); |
| 125 | std::iota(shift_transpose_dimensions.begin(), |
| 126 | shift_transpose_dimensions.end(), 0); |
| 127 | shift_transpose_dimensions.insert(shift_transpose_dimensions.begin(), 1, |
| 128 | shape.dimensions_size()); |
| 129 | |
| 130 | // Shift the input by sizeof(T) bytes and apply bit_mask to unpack. |
| 131 | XlaOp shifted_input = ShiftRightLogical( |
| 132 | broadcast_input, Transpose(Broadcast(shift_bits, shape.dimensions()), |
| 133 | shift_transpose_dimensions)); |
| 134 | XlaOp unpack_input = |
| 135 | And(shifted_input, xla::ConstantR0<uint32>(builder, bit_mask)); |
| 136 | |
| 137 | XlaOp result; |
| 138 | |
| 139 | if (mode_string == "MIN_COMBINED") { |
nothing calls this directly
no test coverage detected