def quantize(self, z):
        """Quantizes z, returns quantized zhat, same shape as z."""
        quantized = round_ste(self.bound(z))
        half_width = self._levels // 2  # Renormalize to [-1, 1].
        return quantized / half_width
Quantizes z, returns quantized zhat, same shape as z.
mlx_audio/tts/models/spark/modules/finite_scalar_quantization.py