def testNoBatchStatic(self): x = np.array([[1., 0], [2, 1]]) # np.linalg.cholesky(y) y = np.array([[1., 2], [2, 5]]) # np.matmul(x, x.T) y_actual = tfb.CholeskyOuterProduct().forward(x=x) x_actual = tfb.CholeskyOuterProduct().inverse(y=y) [y_actual_, x_actual_] = self.evaluate([y_actual, x_actual]) self.assertAllEqual([2, 2], y_actual.shape) self.assertAllEqual([2, 2], x_actual.shape) self.assertAllClose(y, y_actual_) self.assertAllClose(x, x_actual_)
def testNoBatchStatic(self):
x = np.array([[1., 0], [2, 1]]) # np.linalg.cholesky(y)
y = np.array([[1., 2], [2, 5]]) # np.matmul(x, x.T)
with self.test_session() as sess:
y_actual = tfb.CholeskyOuterProduct().forward(x=x)
x_actual = tfb.CholeskyOuterProduct().inverse(y=y)
[y_actual_, x_actual_] = sess.run([y_actual, x_actual])
self.assertAllEqual([2, 2], y_actual.get_shape())
self.assertAllEqual([2, 2], x_actual.get_shape())
self.assertAllClose(y, y_actual_)
self.assertAllClose(x, x_actual_)
def testNoBatchDeferred(self):
x_ = np.array([[1., 0], [2, 1]]) # np.linalg.cholesky(y)
y_ = np.array([[1., 2], [2, 5]]) # np.matmul(x, x.T)
x = tf.compat.v1.placeholder_with_default(x_, shape=None)
y = tf.compat.v1.placeholder_with_default(y_, shape=None)
y_actual = tfb.CholeskyOuterProduct().forward(x=x)
x_actual = tfb.CholeskyOuterProduct().inverse(y=y)
[y_actual_, x_actual_] = self.evaluate([y_actual, x_actual])
# Shapes are always known in eager.
if not tf.executing_eagerly():
self.assertEqual(None, y_actual.shape)
self.assertEqual(None, x_actual.shape)
self.assertAllClose(y_, y_actual_)
self.assertAllClose(x_, x_actual_)
def testNoBatchDeferred(self):
x = np.array([[1., 0], [2, 1]]) # np.linalg.cholesky(y)
y = np.array([[1., 2], [2, 5]]) # np.matmul(x, x.T)
with self.test_session() as sess:
x_pl = tf.placeholder(tf.float32)
y_pl = tf.placeholder(tf.float32)
y_actual = tfb.CholeskyOuterProduct().forward(x=x_pl)
x_actual = tfb.CholeskyOuterProduct().inverse(y=y_pl)
[y_actual_, x_actual_] = sess.run([y_actual, x_actual],
feed_dict={x_pl: x, y_pl: y})
self.assertEqual(None, y_actual.get_shape())
self.assertEqual(None, x_actual.get_shape())
self.assertAllClose(y, y_actual_)
self.assertAllClose(x, x_actual_)
def testBijectorMatrix(self):
bijector = tfb.CholeskyOuterProduct(validate_args=True)
self.assertEqual("cholesky_outer_product", bijector.name)
x = [[[1., 0], [2, 1]], [[np.sqrt(2.), 0], [np.sqrt(8.), 1]]]
y = np.matmul(x, np.transpose(x, axes=(0, 2, 1)))
# Fairly easy to compute differentials since we have 2x2.
dx_dy = [[[2. * 1, 0, 0],
[2, 1, 0],
[0, 2 * 2, 2 * 1]],
[[2 * np.sqrt(2.), 0, 0],
[np.sqrt(8.), np.sqrt(2.), 0],
[0, 2 * np.sqrt(8.), 2 * 1]]]
ildj = -np.sum(
np.log(np.asarray(dx_dy).diagonal(
offset=0, axis1=1, axis2=2)),
axis=1)
self.assertAllEqual((2, 2, 2), bijector.forward(x).shape)
self.assertAllEqual((2, 2, 2), bijector.inverse(y).shape)
self.assertAllClose(y, self.evaluate(bijector.forward(x)))
self.assertAllClose(x, self.evaluate(bijector.inverse(y)))
self.assertAllClose(
ildj,
self.evaluate(
bijector.inverse_log_det_jacobian(
y, event_ndims=2)), atol=0., rtol=1e-7)
self.assertAllClose(
self.evaluate(-bijector.inverse_log_det_jacobian(
y, event_ndims=2)),
self.evaluate(bijector.forward_log_det_jacobian(
x, event_ndims=2)),
atol=0.,
rtol=1e-7)
def __init__(self,
base_kernel,
fixed_inputs,
diag_shift=None,
validate_args=False,
name='SchurComplement'):
"""Construct a SchurComplement kernel instance.
Args:
base_kernel: A `PositiveSemidefiniteKernel` instance, the kernel used to
build the block matrices of which this kernel computes the Schur
complement.
fixed_inputs: A Tensor, representing a collection of inputs. The Schur
complement that this kernel computes comes from a block matrix, whose
bottom-right corner is derived from `base_kernel.matrix(fixed_inputs,
fixed_inputs)`, and whose top-right and bottom-left pieces are
constructed by computing the base_kernel at pairs of input locations
together with these `fixed_inputs`. `fixed_inputs` is allowed to be an
empty collection (either `None` or having a zero shape entry), in which
case the kernel falls back to the trivial application of `base_kernel`
to inputs. See class-level docstring for more details on the exact
computation this does; `fixed_inputs` correspond to the `Z` structure
discussed there. `fixed_inputs` is assumed to have shape `[b1, ..., bB,
N, f1, ..., fF]` where the `b`'s are batch shape entries, the `f`'s are
feature_shape entries, and `N` is the number of fixed inputs. Use of
this kernel entails a 1-time O(N^3) cost of computing the Cholesky
decomposition of the k(Z, Z) matrix. The batch shape elements of
`fixed_inputs` must be broadcast compatible with
`base_kernel.batch_shape`.
diag_shift: A floating point scalar to be added to the diagonal of the
divisor_matrix before computing its Cholesky.
validate_args: If `True`, parameters are checked for validity despite
possibly degrading runtime performance.
Default value: `False`
name: Python `str` name prefixed to Ops created by this class.
Default value: `"SchurComplement"`
"""
with tf.compat.v1.name_scope(name, values=[base_kernel,
fixed_inputs]) as name:
dtype = dtype_util.common_dtype([base_kernel, fixed_inputs],
tf.float32)
self._base_kernel = base_kernel
self._fixed_inputs = (None if fixed_inputs is None else
tf.convert_to_tensor(value=fixed_inputs,
dtype=dtype))
if not self._is_empty_fixed_inputs():
# We create and store this matrix here, so that we get the caching
# benefit when we later access its cholesky. If we computed the matrix
# every time we needed the cholesky, the bijector cache wouldn't be hit.
self._divisor_matrix = base_kernel.matrix(
fixed_inputs, fixed_inputs)
if diag_shift is not None:
self._divisor_matrix = _add_diagonal_shift(
self._divisor_matrix, diag_shift)
self._cholesky_bijector = tfb.Invert(tfb.CholeskyOuterProduct())
super(SchurComplement, self).__init__(base_kernel.feature_ndims,
dtype=dtype,
name=name)
def testNoBatchStaticJacobian(self):
x = np.eye(2)
bijector = tfb.CholeskyOuterProduct()
# The Jacobian matrix is 2 * tf.eye(2), which has jacobian determinant 4.
self.assertAllClose(
np.log(4),
self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=2)))
def testNoBatchDynamicJacobian(self):
bijector = tfb.CholeskyOuterProduct()
x = tf1.placeholder_with_default(np.eye(2, dtype=np.float32),
shape=None)
log_det_jacobian = bijector.forward_log_det_jacobian(x, event_ndims=2)
# The Jacobian matrix is 2 * tf.eye(2), which has jacobian determinant 4.
self.assertAllClose(np.log(4), self.evaluate(log_det_jacobian))
def testNoBatchDynamicJacobian(self):
x = np.eye(2)
bijector = tfb.CholeskyOuterProduct()
x_pl = tf.placeholder(tf.float32)
with self.test_session():
log_det_jacobian = bijector.forward_log_det_jacobian(x_pl,
event_ndims=2)
# The Jacobian matrix is 2 * tf.eye(2), which has jacobian determinant 4.
self.assertAllClose(np.log(4), log_det_jacobian.eval({x_pl: x}))
def testCholeskyFn(self):
def robust_cholesky(x):
return tf.linalg.cholesky(
tf.linalg.set_diag(x,
tf.linalg.diag_part(x) + 1.))
bijector = tfb.CholeskyOuterProduct(cholesky_fn=robust_cholesky)
# We'll add one to the diagonal, so we'll expect the inverse to be
# the `sqrt(diagonal + 1)`.
x = 3 * np.eye(3)
self.assertAllClose(2. * np.eye(3), self.evaluate(bijector.inverse(x)))
