def __call__(self, z):
        """
        einstein notation
        b - batch
        n - sequence (or flattened spatial dimensions)
        d - feature dimension
        c - number of codebook dim
        """

        is_img_or_video = z.ndim >= 4
        need_move_channel_last = is_img_or_video or self.channel_first

        # standardize image or video into (batch, seq, dimension)

        if need_move_channel_last:
            z = mx.reshape(z, (z.shape[0], -1, z.shape[-1]))
            # z = mx.reshape(z, (z.shape[0], -1))

        assert (
            z.shape[-1] == self.dim
        ), f"expected dimension of {self.dim} but found dimension of {z.shape[-1]}"

        z = self.project_in(z)

        z = mx.reshape(z, (z.shape[0], z.shape[1], self.num_codebooks, z.shape[-1]))

        # whether to force quantization step to be full precision or not

        force_f32 = self.force_quantization_f32

        orig_dtype = z.dtype

        if force_f32 and orig_dtype not in self.allowed_dtypes:
            z = z.float()

        codes = self.quantize(z)

        # returning indices could be optional

        indices = None

        if self.return_indices:
            indices = self.codes_to_indices(codes)

        codes = mx.reshape(codes, (codes.shape[0], codes.shape[1], -1))

        codes = codes.astype(orig_dtype)

        # project out

        out = self.project_out(codes)

        # reconstitute image or video dimensions

        if need_move_channel_last:
            out = mx.reshape(out, (out.shape[0], -1, out.shape[-1]))
            out = mx.reshape(
                out, (out.shape[0], out.shape[1], out.shape[2], out.shape[-1])
            )

            indices = mx.reshape(indices, (indices.shape[0], -1, indices.shape[-1]))

        if not self.keep_num_codebooks_dim and self.return_indices:
            indices = mx.reshape(indices, (indices.shape[0], -1, indices.shape[-1]))

        # return quantized output and indices

        return out, indices

        einstein notation
        b - batch
        n - sequence (or flattened spatial dimensions)
        d - feature dimension
        c - number of codebook dim
        
mlx_audio/tts/models/spark/modules/finite_scalar_quantization.py