def __call__(self, x):
        """
        x: a 3-dimensional tensor in tdnn-based architecture (B,F,T)
            or a 4-dimensional tensor in resnet architecture (B,C,F,T)
            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
        """
        if len(x.shape) == 4:
            x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])
        assert len(x.shape) == 3

        if self.global_context_att:
            context_mean = mx.mean(x, axis=-1)[:, :, None]
            context_mean = mx.broadcast_to(context_mean, x.shape)
            context_std = mx.sqrt(mx.var(x, axis=-1) + 1e-7)[:, :, None]
            context_std = mx.broadcast_to(context_std, x.shape)
            x_in = mx.concatenate((x, context_mean, context_std), axis=1)
        else:
            x_in = x

        # DON'T use ReLU here! ReLU may be hard to converge.
        alpha = mx.tanh(
            self.linear1(x_in.transpose(0, 2, 1)).transpose(0, 2, 1)
        )  # alpha = F.relu(self.linear1(x_in))
        alpha = mx.softmax(
            self.linear2(alpha.transpose(0, 2, 1)).transpose(0, 2, 1), axis=2
        )
        mean = mx.sum(alpha * x, axis=2)
        var = mx.sum(alpha * (x**2), axis=2) - mean**2
        std = mx.sqrt(mx.clip(var, 1e-7, None))
        return mx.concatenate([mean, std], axis=1)

        x: a 3-dimensional tensor in tdnn-based architecture (B,F,T)
            or a 4-dimensional tensor in resnet architecture (B,C,F,T)
            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
        
mlx_audio/tts/models/spark/modules/speaker/pooling_layers.py