def testNoBatchMultivariateFullDynamic(self):
with self.cached_session() as sess:
x = array_ops.placeholder(dtypes.float32, name="x")
mu = array_ops.placeholder(dtypes.float32, name="mu")
scale_diag = array_ops.placeholder(dtypes.float32,
name="scale_diag")
x_value = np.array([[1., 1]], dtype=np.float32)
mu_value = np.array([1., -1], dtype=np.float32)
scale_diag_value = np.array([2., 2], dtype=np.float32)
feed_dict = {
x: x_value,
mu: mu_value,
scale_diag: scale_diag_value,
}
bijector = Affine(shift=mu, scale_diag=scale_diag)
self.assertAllClose([[3., 1]],
sess.run(bijector.forward(x), feed_dict))
self.assertAllClose([[0., 1]],
sess.run(bijector.inverse(x), feed_dict))
self.assertAllClose(
-np.log(4),
sess.run(bijector.inverse_log_det_jacobian(x, event_ndims=1),
feed_dict))
def testCompareToBijector(self):
"""Demonstrates equivalence between TD, Bijector approach and AR dist."""
sample_shape = np.int32([4, 5])
batch_shape = np.int32([])
event_size = np.int32(2)
with self.cached_session() as sess:
batch_event_shape = np.concatenate([batch_shape, [event_size]],
axis=0)
sample0 = array_ops.zeros(batch_event_shape)
affine = Affine(scale_tril=self._random_scale_tril(event_size))
ar = autoregressive_lib.Autoregressive(self._normal_fn(affine),
sample0,
validate_args=True)
ar_flow = MaskedAutoregressiveFlow(is_constant_jacobian=True,
shift_and_log_scale_fn=lambda x:
[None, affine.forward(x)],
validate_args=True)
td = transformed_distribution_lib.TransformedDistribution(
distribution=normal_lib.Normal(loc=0., scale=1.),
bijector=ar_flow,
event_shape=[event_size],
batch_shape=batch_shape,
validate_args=True)
x_shape = np.concatenate([sample_shape, batch_shape, [event_size]],
axis=0)
x = 2. * self._rng.random_sample(x_shape).astype(np.float32) - 1.
td_log_prob_, ar_log_prob_ = sess.run(
[td.log_prob(x), ar.log_prob(x)])
self.assertAllClose(td_log_prob_, ar_log_prob_, atol=0., rtol=1e-6)
def testIdentityWithVDVTUpdate(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# Corresponds to scale = [[10, 0, 0], [0, 2, 0], [0, 0, 3]]
bijector = Affine(shift=mu,
scale_identity_multiplier=2.,
scale_perturb_diag=[2., 1],
scale_perturb_factor=[[2., 0], [0., 0],
[0, 1]])
bijector_ref = Affine(shift=mu, scale_diag=[10., 2, 3])
x = [1., 2, 3] # Vector.
self.assertAllClose([9., 3, 8], run(bijector.forward, x))
self.assertAllClose(run(bijector_ref.forward, x),
run(bijector.forward, x))
self.assertAllClose([0.2, 1.5, 4 / 3.],
run(bijector.inverse, x))
self.assertAllClose(run(bijector_ref.inverse, x),
run(bijector.inverse, x))
self.assertAllClose(-np.log(60.),
run(bijector.inverse_log_det_jacobian, x))
self.assertAllClose(
run(bijector.inverse_log_det_jacobian, x),
run(bijector_ref.inverse_log_det_jacobian, x))
def testBatchMultivariateFullDynamic(self):
with self.test_session() as sess:
x = array_ops.placeholder(dtypes.float32, name="x")
mu = array_ops.placeholder(dtypes.float32, name="mu")
scale_diag = array_ops.placeholder(dtypes.float32,
name="scale_diag")
event_ndims = array_ops.placeholder(dtypes.int32,
name="event_ndims")
x_value = np.array([[[1., 1]]], dtype=np.float32)
mu_value = np.array([[1., -1]], dtype=np.float32)
scale_diag_value = np.array([[2., 2]], dtype=np.float32)
event_ndims_value = 1
feed_dict = {
x: x_value,
mu: mu_value,
scale_diag: scale_diag_value,
event_ndims: event_ndims_value
}
bijector = Affine(shift=mu,
scale_diag=scale_diag,
event_ndims=event_ndims)
self.assertEqual(1, sess.run(bijector.event_ndims, feed_dict))
self.assertAllClose([[[3., 1]]],
sess.run(bijector.forward(x), feed_dict))
self.assertAllClose([[[0., 1]]],
sess.run(bijector.inverse(x), feed_dict))
self.assertAllClose([-np.log(4)],
sess.run(bijector.inverse_log_det_jacobian(x),
feed_dict))
def testDiagWithVDVTUpdate(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# Corresponds to scale = [[10, 0, 0], [0, 3, 0], [0, 0, 5]]
bijector = Affine(shift=mu,
scale_diag=[2., 3, 4],
scale_perturb_diag=[2., 1],
scale_perturb_factor=[[2., 0], [0., 0],
[0, 1]])
bijector_ref = Affine(shift=mu, scale_diag=[10., 3, 5])
self.assertEqual(1, bijector.event_ndims.eval()) # "is vector"
x = [1., 2, 3] # Vector.
self.assertAllClose([9., 5, 14], run(bijector.forward, x))
self.assertAllClose(run(bijector_ref.forward, x),
run(bijector.forward, x))
self.assertAllClose([0.2, 1., 0.8], run(bijector.inverse, x))
self.assertAllClose(run(bijector_ref.inverse, x),
run(bijector.inverse, x))
self.assertAllClose(-np.log(150.),
run(bijector.inverse_log_det_jacobian, x))
self.assertAllClose(
run(bijector.inverse_log_det_jacobian, x),
run(bijector_ref.inverse_log_det_jacobian, x))
def testTriLWithVDVTUpdateNoDiagonal(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# Corresponds to scale = [[6, 0, 0], [1, 3, 0], [2, 3, 5]]
bijector = Affine(
shift=mu,
scale_tril=[[2., 0, 0], [1, 3, 0], [2, 3, 4]],
scale_perturb_diag=None,
scale_perturb_factor=[[2., 0], [0., 0], [0, 1]])
bijector_ref = Affine(
shift=mu, scale_tril=[[6., 0, 0], [1, 3, 0], [2, 3, 5]])
self.assertEqual(1, bijector.event_ndims.eval()) # "is vector"
x = [1., 2, 3] # Vector.
self.assertAllClose([5., 6, 22], run(bijector.forward, x))
self.assertAllClose(
run(bijector_ref.forward, x), run(bijector.forward, x))
self.assertAllClose([1 / 3., 8 / 9., 4 / 30.], run(bijector.inverse, x))
self.assertAllClose(
run(bijector_ref.inverse, x), run(bijector.inverse, x))
self.assertAllClose(-np.log(90.),
run(bijector.inverse_log_det_jacobian, x))
self.assertAllClose(
run(bijector.inverse_log_det_jacobian, x),
run(bijector_ref.inverse_log_det_jacobian, x))
def testCompareToBijector(self):
"""Demonstrates equivalence between TD, Bijector approach and AR dist."""
sample_shape = np.int32([4, 5])
batch_shape = np.int32([])
event_size = np.int32(2)
with self.test_session() as sess:
batch_event_shape = np.concatenate([batch_shape, [event_size]], axis=0)
sample0 = array_ops.zeros(batch_event_shape)
affine = Affine(scale_tril=self._random_scale_tril(event_size))
ar = autoregressive_lib.Autoregressive(
self._normal_fn(affine), sample0, validate_args=True)
ar_flow = MaskedAutoregressiveFlow(
is_constant_jacobian=True,
shift_and_log_scale_fn=lambda x: [None, affine.forward(x)],
validate_args=True)
td = transformed_distribution_lib.TransformedDistribution(
distribution=normal_lib.Normal(loc=0., scale=1.),
bijector=ar_flow,
event_shape=[event_size],
batch_shape=batch_shape,
validate_args=True)
x_shape = np.concatenate(
[sample_shape, batch_shape, [event_size]], axis=0)
x = 2. * self._rng.random_sample(x_shape).astype(np.float32) - 1.
td_log_prob_, ar_log_prob_ = sess.run([td.log_prob(x), ar.log_prob(x)])
self.assertAllClose(td_log_prob_, ar_log_prob_, atol=0., rtol=1e-6)
def testBatchMultivariateFullDynamic(self):
with self.test_session() as sess:
x = array_ops.placeholder(dtypes.float32, name="x")
mu = array_ops.placeholder(dtypes.float32, name="mu")
scale_diag = array_ops.placeholder(dtypes.float32, name="scale_diag")
event_ndims = array_ops.placeholder(dtypes.int32, name="event_ndims")
x_value = np.array([[[1., 1]]], dtype=np.float32)
mu_value = np.array([[1., -1]], dtype=np.float32)
scale_diag_value = np.array([[2., 2]], dtype=np.float32)
event_ndims_value = 1
feed_dict = {
x: x_value,
mu: mu_value,
scale_diag: scale_diag_value,
event_ndims: event_ndims_value
}
bijector = Affine(
shift=mu, scale_diag=scale_diag, event_ndims=event_ndims)
self.assertEqual(1, sess.run(bijector.event_ndims, feed_dict))
self.assertAllClose([[[3., 1]]], sess.run(bijector.forward(x), feed_dict))
self.assertAllClose([[[0., 1]]], sess.run(bijector.inverse(x), feed_dict))
self.assertAllClose([-np.log(4)],
sess.run(
bijector.inverse_log_det_jacobian(x), feed_dict))
def testScalarEventIdentityScale(self):
with self.test_session() as sess:
doubler = Affine(
scale_identity_multiplier=2.,
event_ndims=0)
doubler2 = doubler.inverse_log_det_jacobian(2.)
doubler2_ildj_ = sess.run([doubler2])
self.assertAllClose([-np.log(2.)], doubler2_ildj_)
def testEventNdimsLargerThanOneRaises(self):
with self.test_session():
mu = [1., -1]
with self.assertRaisesRegexp(
ValueError, (r"event_ndims\(2\) was not 0 or 1")):
# Scale corresponds to 2x2 identity matrix.
bijector = Affine(shift=mu, event_ndims=2, validate_args=True)
bijector.forward([1., 1.]).eval()
def testNoBatchMultivariateRaisesWhenSingular(self):
with self.test_session():
mu = [1., -1]
bijector = Affine(
shift=mu,
# Has zero on the diagonal.
scale_diag=[0., 1],
validate_args=True)
with self.assertRaisesOpError("Condition x > 0"):
bijector.forward([1., 1.]).eval()
def testNoBatchMultivariateRaisesWhenSingular(self):
with self.test_session():
mu = [1., -1]
bijector = Affine(
shift=mu,
# Has zero on the diagonal.
scale_diag=[0., 1],
validate_args=True)
with self.assertRaisesOpError("diagonal part must be non-zero"):
bijector.forward([1., 1.]).eval()
def testScaleZeroScalarRaises(self):
with self.test_session():
mu = -1.
# Check Identity matrix with zero scaling.
bijector = Affine(
shift=mu,
scale_identity_multiplier=0.,
event_ndims=0,
validate_args=True)
with self.assertRaisesOpError("identity_multiplier should be non-zero"):
bijector.forward(1.).eval()
def _testLegalInputs(self, shift=None, scale_params=None, x=None):
def _powerset(x):
s = list(x)
return itertools.chain.from_iterable(
itertools.combinations(s, r) for r in range(len(s) + 1))
for args in _powerset(scale_params.items()):
with self.test_session():
args = dict(args)
scale_args = dict({"x": x}, **args)
scale = self._makeScale(**scale_args)
bijector_args = dict({"event_ndims": 1}, **args)
# We haven't specified enough information for the scale.
if scale is None:
with self.assertRaisesRegexp(ValueError, ("must be specified.")):
bijector = Affine(shift=shift, **bijector_args)
else:
bijector = Affine(shift=shift, **bijector_args)
np_x = x
# For the case a vector is passed in, we need to make the shape
# match the matrix for matmul to work.
if x.ndim == scale.ndim - 1:
np_x = np.expand_dims(x, axis=-1)
forward = np.matmul(scale, np_x) + shift
if x.ndim == scale.ndim - 1:
forward = np.squeeze(forward, axis=-1)
self.assertAllClose(forward, bijector.forward(x).eval())
backward = np.linalg.solve(scale, np_x - shift)
if x.ndim == scale.ndim - 1:
backward = np.squeeze(backward, axis=-1)
self.assertAllClose(backward, bijector.inverse(x).eval())
ildj = -np.log(np.abs(np.linalg.det(scale)))
# TODO(jvdillon): We need to make it so the scale_identity_multiplier
# case does not deviate in expected shape. Fixing this will get rid of
# these special cases.
if (ildj.ndim > 0 and (len(scale_args) == 1 or (
len(scale_args) == 2 and
scale_args.get("scale_identity_multiplier", None) is not None))):
ildj = np.squeeze(ildj[0])
elif ildj.ndim < scale.ndim - 2:
ildj = np.reshape(ildj, scale.shape[0:-2])
self.assertAllClose(ildj, bijector.inverse_log_det_jacobian(x).eval())
def testMinEventNdimsShapeChangingRemoveDims(self):
chain = Chain([ShapeChanging(3, 0)])
self.assertEqual(3, chain.forward_min_event_ndims)
self.assertEqual(0, chain.inverse_min_event_ndims)
chain = Chain([ShapeChanging(3, 0), Affine()])
self.assertEqual(3, chain.forward_min_event_ndims)
self.assertEqual(0, chain.inverse_min_event_ndims)
chain = Chain([Affine(), ShapeChanging(3, 0)])
self.assertEqual(4, chain.forward_min_event_ndims)
self.assertEqual(1, chain.inverse_min_event_ndims)
chain = Chain([ShapeChanging(3, 0), ShapeChanging(3, 0)])
self.assertEqual(6, chain.forward_min_event_ndims)
self.assertEqual(0, chain.inverse_min_event_ndims)
def testNoBatchMultivariateIdentity(self):
with self.test_session() as sess:
placeholder = array_ops.placeholder(dtypes.float32, name="x")
def static_run(fun, x, **kwargs):
return fun(x, **kwargs).eval()
def dynamic_run(fun, x_value, **kwargs):
x_value = np.array(x_value)
return sess.run(fun(placeholder, **kwargs),
feed_dict={placeholder: x_value})
for run in (static_run, dynamic_run):
mu = [1., -1]
# Multivariate
# Corresponds to scale = [[1., 0], [0, 1.]]
bijector = Affine(shift=mu)
x = [1., 1]
# matmul(sigma, x) + shift
# = [-1, -1] + [1, -1]
self.assertAllClose([2., 0], run(bijector.forward, x))
self.assertAllClose([0., 2], run(bijector.inverse, x))
# x is a 2-batch of 2-vectors.
# The first vector is [1, 1], the second is [-1, -1].
# Each undergoes matmul(sigma, x) + shift.
x = [[1., 1], [-1., -1]]
self.assertAllClose([[2., 0], [0., -2]],
run(bijector.forward, x))
self.assertAllClose([[0., 2], [-2., 0]],
run(bijector.inverse, x))
self.assertAllClose(
0., run(bijector.inverse_log_det_jacobian,
x,
event_ndims=1))
def testIdentityWithTriL(self):
with self.cached_session() as sess:
placeholder = array_ops.placeholder(dtypes.float32, name="x")
def static_run(fun, x, **kwargs):
return fun(x, **kwargs).eval()
def dynamic_run(fun, x_value, **kwargs):
x_value = np.array(x_value)
return sess.run(
fun(placeholder, **kwargs), feed_dict={placeholder: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# scale = [[2., 0], [2, 2]]
bijector = Affine(
shift=mu,
scale_identity_multiplier=1.,
scale_tril=[[1., 0], [2., 1]])
x = [[1., 2]] # One multivariate sample.
self.assertAllClose([[1., 5]], run(bijector.forward, x))
self.assertAllClose([[1., 0.5]], run(bijector.inverse, x))
self.assertAllClose(
-np.log(4.),
run(bijector.inverse_log_det_jacobian, x, event_ndims=1))
def testScaleDiagAndEventNdimsZeroRaises(self):
# Check Diag matrix with zero scaling.
with self.assertRaisesRegexp(ValueError, "only scale argument"):
Affine(shift=None,
scale_diag=[0.0],
event_ndims=0,
validate_args=True)
def testWeirdSampleNoBatchScalarViaIdentity(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# Corresponds to scale = 2.
bijector = Affine(shift=mu,
scale_identity_multiplier=2.,
event_ndims=0)
self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar"
x = [[1., 2, 3], [4, 5, 6]] # Weird sample shape.
self.assertAllClose([[1., 3, 5], [7, 9, 11]],
run(bijector.forward, x))
self.assertAllClose([[1., 1.5, 2.], [2.5, 3, 3.5]],
run(bijector.inverse, x))
self.assertAllClose(-np.log(2.),
run(bijector.inverse_log_det_jacobian, x))
def testNoBatchMultivariateDiag(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = [1., -1]
# Multivariate
# Corresponds to scale = [[2., 0], [0, 1.]]
bijector = Affine(shift=mu, scale_diag=[2., 1])
self.assertEqual(1, bijector.event_ndims.eval()) # "is vector"
x = [1., 1]
# matmul(sigma, x) + shift
# = [-1, -1] + [1, -1]
self.assertAllClose([3., 0], run(bijector.forward, x))
self.assertAllClose([0., 2], run(bijector.inverse, x))
self.assertAllClose(-np.log(2.),
run(bijector.inverse_log_det_jacobian, x))
# x is a 2-batch of 2-vectors.
# The first vector is [1, 1], the second is [-1, -1].
# Each undergoes matmul(sigma, x) + shift.
x = [[1., 1], [-1., -1]]
self.assertAllClose([[3., 0], [-1., -2]],
run(bijector.forward, x))
self.assertAllClose([[0., 2], [-1., 0]],
run(bijector.inverse, x))
self.assertAllClose(-np.log(2.),
run(bijector.inverse_log_det_jacobian, x))
def testIdentityAndDiagWithTriL(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# scale = [[3., 0], [2, 4]]
bijector = Affine(shift=mu,
scale_identity_multiplier=1.0,
scale_diag=[1., 2.],
scale_tril=[[1., 0], [2., 1]])
self.assertEqual(1, bijector.event_ndims.eval()) # "is vector"
x = [[1., 2]] # One multivariate sample.
self.assertAllClose([[2., 9]], run(bijector.forward, x))
self.assertAllClose([[2 / 3., 5 / 12.]],
run(bijector.inverse, x))
self.assertAllClose(-np.log(12.),
run(bijector.inverse_log_det_jacobian, x))
def testIdentityWithDiagUpdate(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# Corresponds to scale = 2
bijector = Affine(
shift=mu,
scale_identity_multiplier=1.,
scale_diag=[1., 1., 1.],
event_ndims=1)
self.assertEqual(1, bijector.event_ndims.eval()) # "is vector"
x = [1., 2, 3] # Three scalar samples (no batches).
self.assertAllClose([1., 3, 5], run(bijector.forward, x))
self.assertAllClose([1., 1.5, 2.], run(bijector.inverse, x))
self.assertAllClose(-np.log(2.**3),
run(bijector.inverse_log_det_jacobian, x))
def testNoBatchMultivariateRaisesWhenSingular(self):
with self.cached_session():
mu = [1., -1]
with self.assertRaisesRegexp(errors.InvalidArgumentError,
"diagonal part must be non-zero"):
_ = Affine(
shift=mu,
# Has zero on the diagonal.
scale_diag=[0., 1],
validate_args=True)
def testTriLWithVDVTUpdate(self):
with self.test_session() as sess:
placeholder = array_ops.placeholder(dtypes.float32, name="x")
def static_run(fun, x, **kwargs):
return fun(x, **kwargs).eval()
def dynamic_run(fun, x_value, **kwargs):
x_value = np.array(x_value)
return sess.run(fun(placeholder, **kwargs),
feed_dict={placeholder: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# Corresponds to scale = [[10, 0, 0], [1, 3, 0], [2, 3, 5]]
bijector = Affine(shift=mu,
scale_tril=[[2., 0, 0], [1, 3, 0], [2, 3,
4]],
scale_perturb_diag=[2., 1],
scale_perturb_factor=[[2., 0], [0., 0],
[0, 1]])
bijector_ref = Affine(shift=mu,
scale_tril=[[10., 0, 0], [1, 3, 0],
[2, 3, 5]])
x = [1., 2, 3] # Vector.
self.assertAllClose([9., 6, 22], run(bijector.forward, x))
self.assertAllClose(run(bijector_ref.forward, x),
run(bijector.forward, x))
self.assertAllClose([0.2, 14 / 15., 4 / 25.],
run(bijector.inverse, x))
self.assertAllClose(run(bijector_ref.inverse, x),
run(bijector.inverse, x))
self.assertAllClose(
-np.log(150.),
run(bijector.inverse_log_det_jacobian, x, event_ndims=1))
self.assertAllClose(
run(bijector.inverse_log_det_jacobian, x, event_ndims=1),
run(bijector_ref.inverse_log_det_jacobian,
x,
event_ndims=1))
def testNoBatchMultivariateFullDynamic(self):
with self.test_session() as sess:
x = array_ops.placeholder(dtypes.float32, name="x")
mu = array_ops.placeholder(dtypes.float32, name="mu")
scale_diag = array_ops.placeholder(dtypes.float32, name="scale_diag")
x_value = np.array([[1., 1]], dtype=np.float32)
mu_value = np.array([1., -1], dtype=np.float32)
scale_diag_value = np.array([2., 2], dtype=np.float32)
feed_dict = {
x: x_value,
mu: mu_value,
scale_diag: scale_diag_value,
}
bijector = Affine(shift=mu, scale_diag=scale_diag)
self.assertAllClose([[3., 1]], sess.run(bijector.forward(x), feed_dict))
self.assertAllClose([[0., 1]], sess.run(bijector.inverse(x), feed_dict))
self.assertAllClose(
-np.log(4),
sess.run(bijector.inverse_log_det_jacobian(x), feed_dict))
def _testLegalInputs(self, shift=None, scale_params=None, x=None):
def _powerset(x):
s = list(x)
return itertools.chain.from_iterable(
itertools.combinations(s, r) for r in range(len(s) + 1))
for args in _powerset(scale_params.items()):
with self.test_session():
args = dict(args)
scale_args = dict({"x": x}, **args)
scale = self._makeScale(**scale_args)
# We haven't specified enough information for the scale.
if scale is None:
with self.assertRaisesRegexp(ValueError,
("must be specified.")):
bijector = Affine(shift=shift, **args)
else:
bijector = Affine(shift=shift, **args)
np_x = x
# For the case a vector is passed in, we need to make the shape
# match the matrix for matmul to work.
if x.ndim == scale.ndim - 1:
np_x = np.expand_dims(x, axis=-1)
forward = np.matmul(scale, np_x) + shift
if x.ndim == scale.ndim - 1:
forward = np.squeeze(forward, axis=-1)
self.assertAllClose(forward, bijector.forward(x).eval())
backward = np.linalg.solve(scale, np_x - shift)
if x.ndim == scale.ndim - 1:
backward = np.squeeze(backward, axis=-1)
self.assertAllClose(backward, bijector.inverse(x).eval())
scale *= np.ones(shape=x.shape[:-1], dtype=scale.dtype)
ildj = -np.log(np.abs(np.linalg.det(scale)))
# TODO (jvdillon): We need to make it so the scale_identity_multiplier id:1102
# https://github.com/imdone/tensorflow/issues/1103
# case does not deviate in expected shape. Fixing this will get rid of
# these special cases.
if (ildj.ndim > 0 and
(len(scale_args) == 1 or
(len(scale_args) == 2 and scale_args.get(
"scale_identity_multiplier", None) is not None))):
ildj = np.squeeze(ildj[0])
elif ildj.ndim < scale.ndim - 2:
ildj = np.reshape(ildj, scale.shape[0:-2])
self.assertAllClose(
ildj,
bijector.inverse_log_det_jacobian(
x, event_ndims=1).eval())
def testChainAffineExp(self):
scale_diag = np.array([1., 2., 3.], dtype=np.float32)
chain = Chain([Affine(scale_diag=scale_diag), Exp()])
x = [0., np.log(2., dtype=np.float32), np.log(3., dtype=np.float32)]
y = [1., 4., 9.]
self.assertAllClose(y, self.evaluate(chain.forward(x)))
self.assertAllClose(x, self.evaluate(chain.inverse(y)))
self.assertAllClose(
np.log(6, dtype=np.float32) + np.sum(x),
self.evaluate(chain.forward_log_det_jacobian(x, event_ndims=1)))
self.assertAllClose(
-np.log(6, dtype=np.float32) - np.sum(x),
self.evaluate(chain.inverse_log_det_jacobian(y, event_ndims=1)))
def testSampleAndLogProbConsistency(self):
batch_shape = []
event_size = 2
with self.cached_session() as sess:
batch_event_shape = np.concatenate([batch_shape, [event_size]],
axis=0)
sample0 = array_ops.zeros(batch_event_shape)
affine = Affine(scale_tril=self._random_scale_tril(event_size))
ar = autoregressive_lib.Autoregressive(self._normal_fn(affine),
sample0,
validate_args=True)
self.run_test_sample_consistent_log_prob(sess.run,
ar,
radius=1.,
center=0.,
rtol=0.01)
def testScaleZeroScalarRaises(self):
with self.test_session():
mu = -1.
# Check Identity matrix with zero scaling.
bijector = Affine(shift=mu,
scale_identity_multiplier=0.0,
event_ndims=0,
validate_args=True)
with self.assertRaisesOpError("Condition x > 0"):
bijector.forward(1.).eval()
# Check Diag matrix with zero scaling.
bijector = Affine(shift=mu,
scale_diag=[0.0],
event_ndims=0,
validate_args=True)
with self.assertRaisesOpError("Condition x > 0"):
bijector.forward(1.).eval()
def testOneBatchScalarViaIdentityIn64BitUserProvidesShiftOnly(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value).astype(np.float64)
x = array_ops.placeholder(dtypes.float64, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = np.float64([1.])
# One batch, scalar.
# Corresponds to scale = 1.
bijector = Affine(shift=mu, event_ndims=0)
self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar"
x = np.float64([1.]) # One sample from one batches.
self.assertAllClose([2.], run(bijector.forward, x))
self.assertAllClose([0.], run(bijector.inverse, x))
self.assertAllClose(0., run(bijector.inverse_log_det_jacobian, x))
def testBatchMultivariateDiag(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value, dtype=np.float32)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = [[1., -1]]
# Corresponds to 1 2x2 matrix, with twos on the diagonal.
scale_diag = [[2., 2]]
bijector = Affine(shift=mu, scale_diag=scale_diag)
x = [[[1., 1]]]
self.assertAllClose([[[3., 1]]], run(bijector.forward, x))
self.assertAllClose([[[0., 1]]], run(bijector.inverse, x))
self.assertAllClose([-np.log(4)],
run(bijector.inverse_log_det_jacobian, x))
def testTwoBatchScalarIdentityViaDiagMultiplier(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = [1., -1]
# Univariate, two batches.
# Corresponds to scale = 1.
bijector = Affine(shift=mu, scale_identity_multiplier=1., event_ndims=0)
self.assertEqual(0, bijector.event_ndims.eval()) # "is scalar"
x = [1., 1] # One sample from each of two batches.
self.assertAllClose([2., 0], run(bijector.forward, x))
self.assertAllClose([0., 2], run(bijector.inverse, x))
self.assertAllClose(0., run(bijector.inverse_log_det_jacobian, x))
def testMinEventNdimsChain(self):
chain = Chain([Exp(), Exp(), Exp()])
self.assertEqual(0, chain.forward_min_event_ndims)
self.assertEqual(0, chain.inverse_min_event_ndims)
chain = Chain([Affine(), Affine(), Affine()])
self.assertEqual(1, chain.forward_min_event_ndims)
self.assertEqual(1, chain.inverse_min_event_ndims)
chain = Chain([Exp(), Affine()])
self.assertEqual(1, chain.forward_min_event_ndims)
self.assertEqual(1, chain.inverse_min_event_ndims)
chain = Chain([Affine(), Exp()])
self.assertEqual(1, chain.forward_min_event_ndims)
self.assertEqual(1, chain.inverse_min_event_ndims)
chain = Chain([Affine(), Exp(), Softplus(), Affine()])
self.assertEqual(1, chain.forward_min_event_ndims)
self.assertEqual(1, chain.inverse_min_event_ndims)
def testDiagWithTriL(self):
with self.test_session() as sess:
def static_run(fun, x):
return fun(x).eval()
def dynamic_run(fun, x_value):
x_value = np.array(x_value)
x = array_ops.placeholder(dtypes.float32, name="x")
return sess.run(fun(x), feed_dict={x: x_value})
for run in (static_run, dynamic_run):
mu = -1.
# scale = [[2., 0], [2, 3]]
bijector = Affine(shift=mu,
scale_diag=[1., 2.],
scale_tril=[[1., 0], [2., 1]])
x = [[1., 2]] # One multivariate sample.
self.assertAllClose([[1., 7]], run(bijector.forward, x))
self.assertAllClose([[1., 1 / 3.]], run(bijector.inverse, x))
self.assertAllClose(-np.log(6.),
run(bijector.inverse_log_det_jacobian, x))
def testScaleZeroScalarRaises(self):
with self.test_session():
mu = -1.
# Check Identity matrix with zero scaling.
bijector = Affine(
shift=mu,
scale_identity_multiplier=0.0,
event_ndims=0,
validate_args=True)
with self.assertRaisesOpError("Condition x > 0"):
bijector.forward(1.).eval()
# Check Diag matrix with zero scaling.
bijector = Affine(
shift=mu, scale_diag=[0.0], event_ndims=0, validate_args=True)
with self.assertRaisesOpError("Condition x > 0"):
bijector.forward(1.).eval()
def testBatchMultivariateIdentity(self):
with self.test_session() as sess:
placeholder = array_ops.placeholder(dtypes.float32, name="x")
def static_run(fun, x, **kwargs):
return fun(x, **kwargs).eval()
def dynamic_run(fun, x_value, **kwargs):
x_value = np.array(x_value)
return sess.run(fun(placeholder, **kwargs),
feed_dict={placeholder: x_value})
for run in (static_run, dynamic_run):
mu = [[1., -1]]
# Corresponds to 1 2x2 matrix, with twos on the diagonal.
scale = 2.
bijector = Affine(shift=mu, scale_identity_multiplier=scale)
x = [[[1., 1]]]
self.assertAllClose([[[3., 1]]], run(bijector.forward, x))
self.assertAllClose([[[0., 1]]], run(bijector.inverse, x))
self.assertAllClose(
-np.log(4),
run(bijector.inverse_log_det_jacobian, x, event_ndims=1))
def testScalePropertyAssertsCorrectly(self):
with self.test_session():
with self.assertRaises(NotImplementedError):
scale = Affine( # pylint:disable=unused-variable
scale_tril=[[1., 0], [2, 1]],
scale_perturb_factor=[2., 1.]).scale
def testEventNdimsLargerThanOneRaises(self):
with self.test_session():
mu = [1., -1]
# Scale corresponds to 2x2 identity matrix.
bijector = Affine(shift=mu, event_ndims=2, validate_args=True)
bijector.forward([1., 1.]).eval()
