Multiclass classification

[syedzayyan]

Hi!

I have been trying to workout multiclass classification and I noticed GPJax does not have any multiclass likelihood like GPFlow. Is there are a preferred GPJax method or could there be a scope for contribution like this:

class SoftmaxLikelihood(gpx.likelihoods.AbstractLikelihood):
    def __init__(self, num_datapoints: int, num_classes: int):
        super().__init__(num_datapoints=num_datapoints)
        self.num_classes = num_classes

    def link_function(self, fx):
        return _SoftmaxBound(fx)

    def log_prob(self, fx, y):
        # fx: (N, C)
        # y:  (N, 1)
        y = y.squeeze(-1)
        log_probs = nn.log_softmax(fx, axis=-1)
        return jnp.sum(log_probs[jnp.arange(len(y)), y])

    def predict(self, dist):
        mu = dist.mean()
        return nn.softmax(mu, axis=-1)


class _SoftmaxBound:
    def __init__(self, fx):
        self.fx = fx

    def log_prob(self, y):
        y = y.squeeze(-1)
        log_probs = nn.log_softmax(self.fx, axis=-1)
        return jnp.sum(log_probs[jnp.arange(len(y)), y])

Also! The slack URL is broken. Is it active anymore?

[thomaspinder]

Yep - would love to see a PR on multiclass classification!

Slack is not officially closed, but I'd prefer to keep discussions open for all here on Github Discussions.