Build a fake quantizer with jax.custom_vjp whose forward pass rounds values onto an integer grid and whose backward pass is a clipped straight-through estimator, then return the quantized values and the gradient that flows back through it.

With the step size scale and the integer range from qmin to qmax, the forward pass maps every entry of x to

where round sends halves to the nearest even integer, as jnp.round does, so 2.5 becomes 2 and 3.5 becomes 4. The true derivative of q is zero almost everywhere, so the backward rule replaces it. Given an upstream gradient c with the same shape as x, called the cotangent, the backward rule returns $c \odot m$, where $\odot$ is elementwise multiplication and the mask m is 1 for entries with $\text{qmin} \le x / \text{scale} \le \text{qmax}$ (both ends included) and 0 for every other entry.

Return a dict with two keys, values holding q(x) and grads holding the gradient for x that the cotangent produces through your custom rule, each a float32 JAX array with the same shape as x.