def __call__(self, input):
        """
        input: a 3-dimensional tensor in xvector architecture
            or a 4-dimensional tensor in resnet architecture
            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
        """
        if len(input.shape) == 4:  # B x F x T
            input = input.reshape(
                input.shape[0], input.shape[1] * input.shape[2], input.shape[3]
            )
        assert len(input.shape) == 3
        bs, f_dim, t_dim = input.shape
        chunks = mx.split(input, self.head_num, axis=1)
        # split
        chunks_out = []
        for i, layer in enumerate(self.heads_att_trans):
            att_score = layer(chunks[i].transpose(0, 2, 1)).transpose(0, 2, 1)
            alpha = mx.softmax(att_score, axis=-1)
            mean = mx.sum(alpha * chunks[i], axis=2)
            var = mx.sum(alpha * chunks[i] ** 2, axis=2) - mean**2
            std = mx.sqrt(mx.clip(var, 1e-7, None))
            chunks_out.append(mx.concatenate((mean, std), axis=1))
        out = mx.concatenate(chunks_out, axis=1)
        return out

        input: a 3-dimensional tensor in xvector architecture
            or a 4-dimensional tensor in resnet architecture
            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
        
mlx_audio/tts/models/spark/modules/speaker/pooling_layers.py