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 _test_fourier_conv2d_common(config, kern, X, Z):
# Use closed-form evaluations as ground truth
Kuu = covariances.Kuu(Z, kern)
Kfu = covariances.Kfu(Z, kern, X)
Kff = kern(X, full_cov=True)
# Test Fourier-feature-based kernel approximator
basis = fourier_basis(kern, num_bases=config.num_bases)
feat_x = basis(X) # [N, B] or [N, L, B]
feat_z = basis(Z)
tol = 3 * config.num_bases ** -0.5
assert allclose(_avg_spatial_inner_product(feat_x, feat_x), Kff, tol, tol)
assert allclose(_avg_spatial_inner_product(feat_x, feat_z), Kfu, tol, tol)
assert allclose(_avg_spatial_inner_product(feat_z, feat_z), Kuu, tol, tol)
del feat_x, feat_z
# Test covariance of functions draw from approximate prior
fx = []
fz = []
count = 0
while count < config.num_samples:
size = min(config.shard_size, config.num_samples - count)
funcs = random_fourier(kern,
basis=basis,
num_bases=config.num_bases,
sample_shape=[size])
fx.append(funcs(X))
fz.append(funcs(Z))
count += size
fx = swap_axes(tf.concat(fx, axis=0), 0, -1) # [L, N, H, W, S]
fz = swap_axes(tf.concat(fz, axis=0), 0, -1) # [L, M, 1, 1, S]
nb = fx.shape.ndims - 4 # num. of batch dimensions
tol += 3 * config.num_samples ** -0.5
frac = 1 / config.num_samples
assert allclose(frac * _avg_spatial_inner_product(fx, fx, nb), Kff, tol, tol)
assert allclose(frac * _avg_spatial_inner_product(fx, fz, nb), Kfu, tol, tol)
assert allclose(frac * _avg_spatial_inner_product(fz, fz, nb), Kuu, tol, tol)
def _Kuf_conv2d_fallback(Z, kernel, X, full_spatial: bool = False, **kwargs):
Kfu = Kfu_dispatch(Z, kernel, X, full_spatial=full_spatial, **kwargs)
ndims_x = X.shape.ndims - 3 # assume x lives in 3d image space
ndims_z = Z.as_images.shape.ndims - 3
if full_spatial:
assert Kfu.shape.ndims == ndims_x + ndims_z + 2
return swap_axes(Kfu, -4, -1) # TODO: this is a hack
# Swap the batch axes of x and z
assert Kfu.shape.ndims == ndims_x + ndims_z
axes = list(range(ndims_x + ndims_z))
perm = axes[ndims_x: ndims_x + ndims_z] + axes[:ndims_x]
return tf.transpose(Kfu, perm)
def _exact_independent(kern: kernels.MultioutputKernel,
Z: TensorLike,
u: TensorLike,
f: TensorLike,
*,
L: TensorLike = None,
diag: TensorLike = None,
basis: AbstractBasis = None,
multioutput_axis: int = 0,
**kwargs):
"""
Return (independent) pathwise updates for each of the latent prior processes
$f$ subject to the condition $p(f | u) = N(f | u, diag)$ on $f = f(Z)$.
"""
u_shape = tuple(u.shape)
f_shape = tuple(f.shape)
assert u_shape[
-1] == kern.num_latent_gps, "Num. outputs != num. latent GPs"
assert u_shape == f_shape[-len(u_shape):], "Incompatible shapes detected"
if basis is None: # finite-dimensional basis used to express the update
basis = kernel_basis(kern, centers=Z)
# 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]
# 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)
if L is None:
if isinstance(Z, inducing_variables.InducingVariables):
K = covariances.Kuu(Z, kern, jitter=0.0)
else:
K = kern(Z, full_cov=True, full_output_cov=False)
K = tf.linalg.set_diag(K, tf.linalg.diag_part(K) + diag[..., 0])
L = tf.linalg.cholesky(K)
# Solve for $Cov(u, u)^{-1}(u - f(Z))$
weights = move_axis(tf.linalg.cholesky_solve(L, err), -1, -3) # [S, L, M]
return MultioutputDenseSampler(basis=basis,
weights=weights,
multioutput_axis=multioutput_axis,
**kwargs)
def convolve(self,
input,
filters,
strides: List = None,
padding: str = None,
dilations: List = None,
data_format: str = None):
if strides is None:
strides = self.strides
if padding is None:
padding = self.padding
if dilations is None:
dilations = self.dilations
if data_format is None:
data_format = self.data_format
shape_out = self.get_shape_out(shape_in=input.shape,
filter_shape=filters.shape,
strides=strides,
dilations=dilations,
data_format=data_format)
_filters = swap_axes(filters, -2, -1)
conv_kwargs = dict(filters=_filters,
padding=padding,
output_shape=shape_out)
if dilations != [1, 1]: # TODO: improve me
assert data_format == 'NHWC'
assert list(strides) == [1, 1]
assert len(dilations) == 2 and dilations[0] == dilations[1]
return tf.nn.atrous_conv2d_transpose(input,
rate=dilations[0],
**conv_kwargs)
return tf.nn.conv2d_transpose(input,
strides=strides,
dilations=dilations,
data_format=data_format,
**conv_kwargs)
def test_cg_conv2d(*args, **kwargs): f = swap_axes(_test_cg_svgp(*args, **kwargs), 0, -1) # [L, N, H, W, S] mf = tf.transpose(tf.reduce_mean(f, [-3, -2, -1])) # [N, L] res = f - tf.reduce_mean(f, axis=-1, keepdims=True) Sff = common.avg_spatial_inner_product(res, batch_dims=f.shape.ndims - 4) return mf, Sff/f.shape[-1] # [L, N, N]
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)