| 67 | } |
| 68 | |
| 69 | xla::XlaOp Quantize(xla::XlaBuilder* b, const xla::XlaOp& input, |
| 70 | const DataType data_type, |
| 71 | const xla::XlaOp& nudged_input_min, |
| 72 | const xla::XlaOp& nudged_input_max, |
| 73 | const xla::XlaOp& input_scale) { |
| 74 | xla::XlaOp one = XlaHelpers::FloatLiteral(b, data_type, 1.0f); |
| 75 | xla::XlaOp inv_scale = xla::Div(one, input_scale); |
| 76 | xla::XlaOp half = XlaHelpers::FloatLiteral(b, data_type, 0.5f); |
| 77 | |
| 78 | xla::XlaOp clamped = xla::Clamp(nudged_input_min, input, nudged_input_max); |
| 79 | xla::XlaOp clamped_shifted = xla::Sub(clamped, nudged_input_min); |
| 80 | xla::XlaOp rounded = |
| 81 | xla::Floor(xla::Add(xla::Mul(clamped_shifted, inv_scale), half)); |
| 82 | return xla::Add(xla::Mul(rounded, input_scale), nudged_input_min); |
| 83 | } |
| 84 | |
| 85 | // Builds a custom_call to a method named 'fake_quant_with_min_max_vars'. |
| 86 | // The method will be provided the input, the min/max range from the original |