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!"