Repository: Blaizzy/mlx-audio
File: mlx_audio/tts/models/spark/modules/residual.py (Python)

Code before:
            "z_q": z_q,
            "indices": indices,
            "dists": dists,
            "vq_loss": vq_loss,
            "perplexity": perplexity,
            "active_num": active_num.astype(mx.float32),
        }

    def vq2emb(self, vq, out_proj=True):
        emb = self.embed_code(vq)
        if out_proj:
            emb = self.out_project(emb)
        return emb

    def tokenize(self, z: mx.array) -> mx.array:
        """tokenize the input tensor"""
        z_e = self.in_project(z.transpose(0, 2, 1)).transpose(0, 2, 1)
        _, indices, _ = self.decode_latents(z_e)
        return indices

    def detokenize(self, indices):
        """detokenize the input indices"""
        # Check if indices are empty
        if indices.shape[0] == 0 or indices.shape[1] == 0:
            # Return an appropriate empty or placeholder tensor
            return mx.zeros((1, self.input_dim, 1))  # Adjust dimensions as needed

        z_q = self.decode_code(indices).transpose(0, 2, 1)
        z_q = self.out_project(z_q)
        return z_q

    def get_emb(self):
        return self.codebook.weight

    def embed_code(self, embed_id):
        return mx.take(self.codebook.weight, embed_id, axis=0)

    def decode_code(self, embed_id):

        return self.embed_code(embed_id).transpose(0, 2, 1)

    def normalize(self, x):
        """Normalize input tensor along dimension 1."""
        norm = mx.sqrt(mx.sum(mx.power(x, 2), axis=1, keepdims=True))
        return x / mx.maximum(norm, 1e-12)

    def decode_latents(self, latents):
        encodings = rearrange(latents, "b d t -> (b t) d")
        codebook = self.codebook.weight

        # L2 normalize encodings and codebook

Review comment (suggestion):
```suggestion
            return mx.zeros((1, self.input_dim, 1))
```

Code after:
            "z_q": z_q,
            "indices": indices,
            "dists": dists,
            "vq_loss": vq_loss,
            "perplexity": perplexity,
            "active_num": active_num.astype(mx.float32),
        }

    def vq2emb(self, vq, out_proj=True):
        emb = self.embed_code(vq)
        if out_proj:
            emb = self.out_project(emb)
        return emb

    def tokenize(self, z: mx.array) -> mx.array:
        """tokenize the input tensor"""
        z_e = self.in_project(z.transpose(0, 2, 1)).transpose(0, 2, 1)
        _, indices, _ = self.decode_latents(z_e)
        return indices

    def detokenize(self, indices):
        """detokenize the input indices"""
        # Check if indices are empty
        if indices.shape[0] == 0 or indices.shape[1] == 0:
            # Return an appropriate empty or placeholder tensor
            return mx.zeros((1, self.input_dim, 1))

        z_q = self.decode_code(indices).transpose(0, 2, 1)
        z_q = self.out_project(z_q)
        return z_q

    def get_emb(self):
        return self.codebook.weight

    def embed_code(self, embed_id):
        return mx.take(self.codebook.weight, embed_id, axis=0)

    def decode_code(self, embed_id):

        return self.embed_code(embed_id).transpose(0, 2, 1)

    def normalize(self, x):
        """Normalize input tensor along dimension 1."""
        norm = mx.sqrt(mx.sum(mx.power(x, 2), axis=1, keepdims=True))
        return x / mx.maximum(norm, 1e-12)

    def decode_latents(self, latents):
        encodings = rearrange(latents, "b d t -> (b t) d")
        codebook = self.codebook.weight

        # L2 normalize encodings and codebook