Generalize likelihoods

[holmrenser]

First thing I would want to get started on is implementing different likelihoods (Poisson, Negative Binomial). It seems that currently only bernoulli is allowed: https://github.com/thomaspinder/GPJax/blob/master/gpjax/likelihoods.py#L131

Having some docs on how to implement different likelihoods would be good as well.

Unfortunately, neither Poisson nor Negative Binomial seem to implemented in distrax at the moment. Is this something that should be done first?

[thomaspinder]

This is very high on the "to do list". If you were interested, then I would be very happy to support and work with you to get some more likelihoods implemented into GPJax. In addition to the Poisson and/or Negative Binomial, I would like to see categorical and exponential likelihood functions supported.

Fortunately, implementing new distributions in Distrax is much more straightforward than in TFP. Further, any TFP distribution can seamlessly be used within distrax e.g.,

d_tfp = tfd.Normal(0., 1.)
d_dx = dx.Normal(1., 2.5)
d_dx.kl_divergence(d_tfp)

My first thought would be that we should open a PR to the Distrax group for any new likelihood that we support as it would be good to minimise the amount of TFP code within GPJax. From the Distrax docs, a new likelihood would simply involve completing the following object

class MyDistribution(distrax.Distribution):

  def __init__(self, ...):
    ...

  def _sample_n(self, key, n):
    samples = ...
    return samples

  def log_prob(self, value):
    log_prob = ...
    return log_prob

  def event_shape(self):
    event_shape = ...
    return event_shape

  def _sample_n_and_log_prob(self, key, n):
    # Optional. Only when more efficient implementation is possible.
    samples, log_prob = ...
    return samples, log_prob

[holmrenser]

I've finally had some time to work on this and opened a PR to distrax for a negative binomial distribution: google-deepmind/distrax#208. Implementation closely follows tensorflow probability.

[patel-zeel]

Given that tfp has most of the distributions, would it make sense to switch to tensorflow_probability.substrates.jax.distributions?

[thomaspinder]

Thanks for flagging this @holmrenser I'll take a look this evening and look towards generalising the likelihoods now I have some time back in my calendar for GPJax related work.

@patel-zeel - TFP is a dependency of Distrax, so I think we can support these distributions if needed. However, moving forward I would love to see Distrax grow as their implementations are significantly cleaner. This, fundamentally, aligns better with the aims of GPJax which is to provide an easilly hackable tool for the research community.

[holmrenser]

Awesome. There might be some work left on my PR, I don't think it currently passes all tests. Would be nice to have another set of eyes :-)

[holmrenser]

Update: I have some working code that uses tensorflow_probability.substrates.jax.distributions. Swapping in distrax distributions once they're available should be easy, but there seems to be little traction on my PR.

@thomaspinder @daniel-dodd shall I open a PR for a negative binomial likelihood based on tfp and update later? Or wait until distrax merges my PR?

[thomaspinder]

Hi @holmrenser. Please open a PR - I'm very happy to be a reviewer for it. I'm comfortable using TFP for now as it is a dependency of Distrax anyway. The benefit to Distrax is the cleanliness of code so it'd be great to switch out the TFP distribution for Distrax once your PR to Distrax is merged, but we shouldn't let this block development of GPJax.

[thomaspinder]

Hey @holmrenser. Is there any update on this PR? I'd love to get more functionality for non-Gaussian likelihoods integrated into GPJax over the coming months.

[holmrenser]

Hi! I'm back from Christmas break and have some code lying around. Short summary: I haven't continued on the distrax PR but implemented a negative binomial likelihood for gpjax using tfp. I'll get my implementation ready for a PR to gpjax this week.

[thomaspinder]

Hi @holmrenser - let me know if you need a hand with this PR. Very keen to have wider likelihood support. If you’re too busy right now then that’s no problem as we can look to open our own PR for the functionality.