def _compute_cholesky_gp(
self,
kernel_matrix: Tensor,
num_data_points: Optional[int] = None,
noise: bool = True,
) -> Tensor:
r"""
Parameters
--------------------
kernel_matrix
Kernel matrix of shape (batch_size, num_data_points, num_data_points).
num_data_points
Number of rows in the kernel_matrix.
noise
Boolean to determine whether to add :math:`\sigma^2I` to the kernel matrix.
This is used in the predictive step if you would like to sample the predictive
covariance matrix without noise. It is set to True in every other case.
Returns
--------------------
Tensor
Cholesky factor :math:`L` of the kernel matrix with added noise :math:`LL^T = K + \sigma^2 I`
of shape (batch_size, num_data_points, num_data_points).
"""
if noise: # Add sigma
kernel_matrix = self.F.broadcast_plus(
kernel_matrix,
self.F.broadcast_mul(
self.sigma**2,
self.F.eye(num_data_points,
ctx=self.ctx,
dtype=self.float_type),
),
)
# Warning: This method is more expensive than the iterative jitter
# but it works for mx.sym
if self.jitter_method == "eig":
return jitter_cholesky_eig(
self.F,
kernel_matrix,
num_data_points,
self.ctx,
self.float_type,
self.diag_weight,
)
elif self.jitter_method == "iter" and self.F is mx.nd:
return jitter_cholesky(
self.F,
kernel_matrix,
num_data_points,
self.ctx,
self.float_type,
self.max_iter_jitter,
self.neg_tol,
self.diag_weight,
self.increase_jitter,
)
else:
return self.F.linalg.potrf(kernel_matrix)
def test_jitter_unit(jitter_method, float_type, ctx=mx.Context("cpu")):
matrix = nd.array([[[1, 2], [3, 4]], [[10, 100], [-21.5, 41]]],
ctx=ctx,
dtype=float_type)
F = mx.nd
num_data_points = matrix.shape[1]
if jitter_method == "eig":
L = jitter_cholesky_eig(F, matrix, num_data_points, ctx, float_type)
elif jitter_method == "iter":
L = jitter_cholesky(F, matrix, num_data_points, ctx, float_type)
assert np.sum(np.isnan(L.asnumpy())) == 0, "NaNs in Cholesky factor!"
def test_jitter_unit(jitter_method, float_type, ctx) -> None:
# TODO: Enable GPU tests on Jenkins
if ctx == mx.Context("gpu") and not check_gpu_support():
return
matrix = nd.array([[[1, 2], [3, 4]], [[10, 100], [-21.5, 41]]],
ctx=ctx,
dtype=float_type)
F = mx.nd
num_data_points = matrix.shape[1]
if jitter_method == "eig":
L = jitter_cholesky_eig(F, matrix, num_data_points, ctx, float_type)
elif jitter_method == "iter":
L = jitter_cholesky(F, matrix, num_data_points, ctx, float_type)
assert np.sum(np.isnan(L.asnumpy())) == 0, "NaNs in Cholesky factor!"
