Skip to content

← JAX Conceptual Deep-Dives step 8 of 15

Medium Primitives

int8 Affine Quantization

Jump to the task

Put values on a fixed grid

Affine quantization represents a real value with an integer code. The scale is the spacing between neighboring reconstructed values, and zero_point is the code that represents real zero.

For this exercise, the codes use the full signed int8 range:

q = clip(round(x / scale) + zero_point, -128, 127)

A smaller scale gives finer steps, but covers a narrower range with the same 256 codes. Values outside that range are clipped to an endpoint. You are given the scale and zero point; choosing them is a separate task.

Worked example

import jax.numpy as jnp

x = jnp.array([-0.75, 0.0, 0.75, 100.0], dtype=jnp.float32)
scale, zero_point = 0.5, 3.0
rounded = jnp.round(x / scale)       # [-2, 0, 2, 200]
shifted = rounded + zero_point      # [1, 3, 5, 203]
q = jnp.clip(shifted, -128.0, 127.0) # [1, 3, 5, 127]

Zero maps to code 3. The two half-way values use ties-to-even rounding: -1.5 rounds to -2, and 1.5 rounds to 2. The last value needs a code larger than 127, so clipping loses the amount by which it exceeded the range.

Problem

Implement affine_quantize(x, scale, zero_point):

  1. Divide x by scale.
  2. Round with jnp.round, which rounds exact ties to the even integer.
  3. Add zero_point to the rounded values.
  4. Clip the codes to [-128.0, 127.0].
  5. Return the floating array of codes with the same shape as x.

x is a one-dimensional float32 array of finite values. scale is a finite positive scalar. zero_point is an integer-valued scalar in [-128, 127], supplied as a float. Empty input returns an empty array.

Keep the codes as floats for this exercise. Do not cast them to int8: this task practices the mapping, not compressed storage.

Details that change the answer

Round before adding the zero point. With an odd zero point, adding first changes which integer is even at a tie: round(0.5) + 1 is 1, whereas round(0.5 + 1) is 2. Clip after adding the zero point so the final codes stay in range.

The full signed int8 range is [-128, 127]. Some quantization schemes reserve a code and use [-127, 127]; that is a different convention, not an int7 format. Use the full range requested here.

Dequantization maps a code back to (q - zero_point) * scale. It cannot recover information lost by rounding or clipping. Keeping these codes as float32 also uses float32 storage; actual int8 storage is a separate step.