def _linear_fallback(Z: TensorLike,
u: TensorLike,
f: TensorLike,
*,
L: TensorLike = None,
diag: TensorLike = None,
basis: AbstractBasis = None,
**kwargs):
u_shape = tuple(u.shape)
f_shape = tuple(f.shape)
assert u_shape[-1] == 1, "Recieved multiple output features"
assert u_shape == f_shape[-len(u_shape):], "Incompatible shapes detected"
# Prepare diagonal term
if diag is None: # used by <GPflow.conditionals>
diag = default_jitter()
if isinstance(diag, float):
diag = tf.convert_to_tensor(diag, dtype=f.dtype)
diag = tf.expand_dims(diag, axis=-1) # [M, 1] or [1, 1] or [1]
# Extract "features" of Z
if basis is None:
if isinstance(Z, inducing_variables.InducingVariables):
feat = inducing_to_tensor(Z) # [M, D]
else:
feat = Z
else:
feat = basis(Z) # [M, D] (maybe a different "D" than above)
# Compute error term and matrix square root $Cov(u, u)^{1/2}$
err = swap_axes(u - f, -3, -1) # [1, M, S]
err -= tf.sqrt(diag) * tf.random.normal(err.shape, dtype=err.dtype)
M, D = feat.shape[-2:]
if L is None:
if D < M:
feat_iDiag = feat * tf.math.reciprocal(diag)
S = tf.matmul(feat_iDiag, feat, transpose_a=True) # [D, D]
L = tf.linalg.cholesky(S + tf.eye(S.shape[-1], dtype=S.dtype))
else:
K = tf.matmul(feat, feat, transpose_b=True) # [M, M]
K = tf.linalg.set_diag(K, tf.linalg.diag_part(K) + diag[..., 0])
L = tf.linalg.cholesky(K)
else:
assert L.shape[-1] == min(M, D) # TODO: improve me
# Solve for $Cov(u, u)^{-1}(u - f(Z))$
if D < M:
feat_iDiag = feat * tf.math.reciprocal(diag)
weights = tf.linalg.adjoint(
tf.linalg.cholesky_solve(
L, tf.matmul(feat_iDiag, err, transpose_a=True)))
else:
iK_err = tf.linalg.cholesky_solve(L, err) # [S, M, 1]
weights = tf.matmul(iK_err, feat, transpose_a=True) # [S, 1, D]
return DenseSampler(basis=basis,
weights=move_axis(weights, -2, -3),
**kwargs)
def __call__(self, x: TensorType, **kwargs) -> tf.Tensor:
self._maybe_initialize(x, **kwargs)
if isinstance(x, InducingVariables): # TODO: Allow this behavior?
x = inducing_to_tensor(x)
proj = tf.tensordot(x, self.weights, axes=[-1, -1]) # [..., B]
feat = tf.cos(proj + self.biases)
return self.output_scale * feat
def initialize(self, x: TensorType, dtype: Any = None):
if isinstance(x, InducingVariables):
x = inducing_to_tensor(x)
if dtype is None:
dtype = x.dtype
self._biases = bias_initializer(self.kernel,
self.num_bases,
dtype=dtype)
self._weights = weight_initializer(self.kernel,
x.shape[-1],
batch_shape=[self.num_bases],
dtype=dtype)
def __call__(self, x: TensorType, multioutput_axis: int = None, **kwargs):
self._maybe_initialize(x, **kwargs)
if isinstance(x, InducingVariables): # TODO: Allow this behavior?
x = inducing_to_tensor(x)
# Compute (batch) tensor dot product
batch_axes = None if (
multioutput_axis is None) else [0, multioutput_axis]
proj = move_axis(
batch_tensordot(self.weights,
x,
axes=[-1, -1],
batch_axes=batch_axes), 1, -1)
ndims = proj.shape.ndims
feat = tf.cos(proj + expand_to(self.biases, axis=1, ndims=ndims))
return expand_to(self.output_scale, axis=1,
ndims=ndims) * feat # [L, N, B]
def initialize(self, x: TensorType, dtype: Any = None):
if isinstance(x, InducingVariables):
x = inducing_to_tensor(x)
if dtype is None:
dtype = x.dtype
biases = []
weights = []
for kernel in self.kernel.latent_kernels:
biases.append(bias_initializer(kernel, self.num_bases,
dtype=dtype))
weights.append(
weight_initializer(kernel,
x.shape[-1],
batch_shape=[self.num_bases],
dtype=dtype))
self._biases = tf.stack(biases, axis=0) # [L, B]
self._weights = tf.stack(weights, axis=0) # [L, B, D]
def _linear_multioutput(Z: inducing_variables.MultioutputInducingVariables,
u: TensorLike,
f: TensorLike,
*,
L: TensorLike = None,
diag: TensorLike = None,
basis: AbstractBasis = None,
multioutput_axis: int = "default",
**kwargs):
assert tuple(u.shape) == tuple(f.shape)
if multioutput_axis == "default":
multioutput_axis = None if (basis is None) else 0
# Prepare diagonal term
if diag is None: # used by <GPflow.conditionals>
diag = default_jitter()
if isinstance(diag, float):
diag = tf.convert_to_tensor(diag, dtype=f.dtype)
diag = tf.expand_dims(diag, axis=-1) # ([L] or []) + ([M] or []) + [1]
# Extract "features" of Z
if basis is None:
if isinstance(Z, inducing_variables.InducingVariables):
feat = inducing_to_tensor(Z) # [L, M, D] or [M, D]
else:
feat = Z
elif isinstance(Z, inducing_variables.SharedIndependentInducingVariables):
feat = basis(Z)
else:
feat = basis(Z,
multioutput_axis=0) # first axis of Z is output-specific
# Compute error term and matrix square root $Cov(u, u)^{1/2}$
err = swap_axes(u - f, -3, -1) # [L, M, S]
err -= tf.sqrt(diag) * tf.random.normal(err.shape, dtype=err.dtype)
M, D = feat.shape[-2:]
if L is None:
if D < M:
feat_iDiag = feat * tf.math.reciprocal(diag)
S = tf.matmul(feat_iDiag, feat,
transpose_a=True) # [L, D, D] or [D, D]
L = tf.linalg.cholesky(S + tf.eye(S.shape[-1], dtype=S.dtype))
else:
K = tf.matmul(feat, feat, transpose_b=True) # [L, M, M] or [M, M]
K = tf.linalg.set_diag(K, tf.linalg.diag_part(K) + diag[..., 0])
L = tf.linalg.cholesky(K)
else:
assert L.shape[-1] == min(M, D) # TODO: improve me
# Solve for $Cov(u, u)^{-1}(u - f(Z))$
if D < M:
feat_iDiag = feat * tf.math.reciprocal(diag)
weights = tf.linalg.adjoint(
tf.linalg.cholesky_solve(
L, tf.matmul(feat_iDiag, err, transpose_a=True)))
else:
iK_err = tf.linalg.cholesky_solve(L, err) # [L, S, M]
weights = tf.matmul(iK_err, feat, transpose_a=True) # [L, S, D]
return MultioutputDenseSampler(
basis=basis,
weights=swap_axes(weights, -3, -2), # [S, L, D]
multioutput_axis=multioutput_axis,
**kwargs)