feat(gpjax/kernels/base.py): add diagonal

[stephen-huan]

Type of changes

• Bug fix
• New feature
• Documentation / docstrings
• Tests
• Other

Checklist

• I've formatted the new code by running poetry run pre-commit run --all-files --show-diff-on-failure before committing.
• I've added tests for new code.
• I've added docstrings for the new code.

Description

AbstractKernel has wrappers for cross_covariance and gram but not for diagonal, when all three are defined by AbstractKernelComputation (diagonal is actually concretely implemented by AbstractKernelComputation).

This PR adds diagonal to AbstractKernel.

Although the change is trivial, I believe it's important for the following reasons.

Currently, the way to get the diagonal of a kernel matrix efficiently would be code like

import jax.numpy as jnp
from gpjax.kernels import DenseKernelComputation, Matern52

kernel = Matern52()
diagonal = DenseKernelComputation().diagonal
# diag([1. 1. 1. 1. 1. 1. 1. 1. 1. 1.])
print(diagonal(kernel, jnp.zeros((10, 3))))

Other than being needlessly verbose, this obscures the fact that most of the *Computation classes (DiagonalKernelComputation, etc.) could be used in place of DenseKernelComputation without changing the underlying implementation from AbstractKernelComputation. Instead, code like

kernel.diagonal(jnp.zeros((10, 3)))

is cleaner and would allow specific kernels to provide a more efficient, specialized implementation (for example, the Matern family has a unit diagonal no matter the input data).

[github-actions]

Thank you for opening your first PR into GPJax!

If you have not heard from us in a while, please feel free to ping @gpjax/developers or anyone who has commented on the PR. Most of our reviewers are volunteers and sometimes things fall through the cracks.

You can also join us on Slack for real-time discussion.

For details on testing, writing docs, and our review process, please see the developer guide

We strive to be a welcoming and open project. Please follow our Code of Conduct.

[thomaspinder]

Hi @stephen-huan - thanks for adding this. Would you be able to include a unit test for this function please?

[stephen-huan]

@thomaspinder I've added tests, sorry for the long delay!

[stephen-huan]

@thomaspinder any progress on this PR?

I added consistency tests that are tangentially related to this PR: now kernel.gram(x) is checked by kernel.cross_covariance(x, x), kernel.diagonal(x) is checked by diagonal(kernel.gram(x)), and kernel.cross_covariance(a, b) is checked by kernel.gram(c) where c is vstack((a, b)).

I also noticed a check of max(a - b) < tol, when it should probably be max(abs(a - b)) < tol.

[thomaspinder]

Thanks @stephen-huan - nice to see this included in GPJax!