def codes_to_indices(self, zhat):
        """Converts a `code` to an index in the codebook."""
        assert zhat.shape[-1] == self.codebook_dim
        zhat = self._scale_and_shift(zhat)
        return (zhat * self._basis).sum(axis=-1).astype(mx.int32)
Converts a `code` to an index in the codebook.
mlx_audio/tts/models/spark/modules/finite_scalar_quantization.py