def fn(residual_fraction):
bijector = tfp.experimental.bijectors.HighwayFlow(
residual_fraction=residual_fraction,
activation_fn=tf.nn.softplus,
bias=tf.zeros(width),
upper_diagonal_weights_matrix=tf.eye(width),
lower_diagonal_weights_matrix=tf.eye(width),
gate_first_n=width)
return tf.reduce_mean(target.log_prob(bijector.forward(x)))
def call(self, x):
bs, c = x.shape
x_t = tf.transpose(x, (1, 0))
m = tf.reduce_mean(x_t, axis=1, keepdims=True)
f = x_t - m
ff_apr = tf.matmul(f, f,
transpose_b=True) / (tf.cast(bs, tf.float32) - 1.0)
ff_apr_shrinked = (1 - self.eps) * ff_apr + tf.eye(c) * self.eps
sqrt = tf.linalg.cholesky(ff_apr_shrinked)
inv_sqrt = tf.linalg.triangular_solve(sqrt, tf.eye(c))
f_hat = tf.matmul(inv_sqrt, f)
decorelated = tf.transpose(f_hat, (1, 0))
return decorelated
def test_weighted_sum_nested_values(self, dtype):
del dtype # not used in this test case.
weights = [0.5, -0.25, -0.25]
states = [(tf.eye(2), tf.ones((2, 2))) for _ in range(3)]
weighted_state_sum = rk_util.weighted_sum(weights, states)
expected_result = (tf.zeros((2, 2)), tf.zeros((2, 2)))
self.assertAllCloseNested(weighted_state_sum, expected_result)
weights = [0.5, -0.25, -0.25, 0]
states = [(tf.eye(2), tf.ones((2, 2))) for _ in range(4)]
weighted_state_sum = rk_util.weighted_sum(weights, states)
expected_result = (tf.zeros((2, 2)), tf.zeros((2, 2)))
self.assertAllCloseNested(weighted_state_sum, expected_result)
def testIllConditioned(self):
# Modified G-S handles ill-conditioned matrices much better (numerically)
# than original G-S.
dim = 200
mat = tf.eye(200, dtype=tf.float64) * 1e-5 + scipy_linalg.hilbert(dim)
mat = tf.math.l2_normalize(mat, axis=-1)
ortho = tfp.math.gram_schmidt(mat)
xtx = tf.matmul(ortho, ortho, transpose_a=True)
self.assertAllClose(
0.,
tf.linalg.norm(tf.eye(dim, dtype=tf.float64) - xtx,
ord='fro',
axis=(-1, -2)))
def eye(N, M=None, k=0, dtype=float): # pylint: disable=invalid-name
"""Returns a 2-D array with ones on the diagonal and zeros elsewhere.
Examples:
```python
eye(2, dtype=int)
-> [[1, 0],
[0, 1]]
eye(2, M=3, dtype=int)
-> [[1, 0, 0],
[0, 1, 0]]
eye(2, M=3, k=1, dtype=int)
-> [[0, 1, 0],
[0, 0, 1]]
eye(3, M=2, k=-1, dtype=int)
-> [[0, 0],
[1, 0],
[0, 1]]
```
Args:
N: integer. Number of rows in output array.
M: Optional integer. Number of cols in output array, defaults to N.
k: Optional integer. Position of the diagonal. The default 0 refers to the
main diagonal. A positive/negative value shifts the diagonal by the
corresponding positions to the right/left.
dtype: Optional, defaults to float. The type of the resulting ndarray.
Could be a python type, a NumPy type or a TensorFlow `DType`.
Returns:
An ndarray with shape (N, M) and requested type.
"""
if dtype:
dtype = utils.to_tf_type(dtype)
if not M:
M = N
if k >= M or -k >= N:
return zeros([N, M], dtype=dtype)
if k:
if k > 0:
result = tf.eye(N, M, dtype=dtype)
zero_cols = tf.zeros([N, abs(k)], dtype=dtype)
result = tf.concat([zero_cols, result], axis=1)
result = tf.slice(result, [0, 0], [N, M])
else:
result = tf.eye(N, M - k, dtype=dtype)
result = tf.slice(result, [0, -k], [N, M])
else:
result = tf.eye(N, M, dtype=dtype)
return utils.tensor_to_ndarray(result)
def testVariableLocation(self):
loc = tf.Variable([[1., 1.]])
scale_row = tf.linalg.LinearOperatorLowerTriangular(
tf.eye(1), is_non_singular=True)
scale_column = tf.linalg.LinearOperatorLowerTriangular(
tf.eye(2), is_non_singular=True)
d = tfd.MatrixNormalLinearOperator(loc,
scale_row,
scale_column,
validate_args=True)
self.evaluate(loc.initializer)
with tf.GradientTape() as tape:
lp = d.log_prob([[0., 0.]])
self.assertIsNotNone(tape.gradient(lp, loc))
def test_subsidiary_bijector_pins(self):
pinned0 = tfde.JointDistributionPinned(
jd_coroutine(), x=1., z=tf.eye(4, batch_shape=[2]))
pinned1 = tfde.JointDistributionPinned(jd_coroutine(), x=1.)
bij0 = pinned0.experimental_default_event_space_bijector()
bij1 = pinned1.experimental_default_event_space_bijector(
z=tf.eye(4, batch_shape=[2]))
self.assertEqual(bij0.inverse_event_shape(pinned0.event_shape),
bij1.inverse_event_shape(pinned0.event_shape))
self.assertEqual(bij0.inverse_event_shape(pinned1.event_shape[:-1]),
bij1.inverse_event_shape(pinned1.event_shape[:-1]))
init = self.evaluate(tf.nest.map_structure(
lambda shp: tf.random.uniform(shp, -2., 2., seed=test_util.test_seed()),
bij0.inverse_event_shape(pinned0.event_shape)))
self.assertAllCloseNested(bij0.forward(init), bij1.forward(init))
def test_log_prob_matches_linear_gaussian_ssm(self):
dim = 2
batch_shape = [3, 1]
seed, *model_seeds = samplers.split_seed(test_util.test_seed(), n=6)
# Sample a random linear Gaussian process.
prior_loc = self.evaluate(
tfd.Normal(0., 1.).sample(batch_shape + [dim],
seed=model_seeds[0]))
prior_scale = self.evaluate(
tfd.InverseGamma(1., 1.).sample(batch_shape + [dim],
seed=model_seeds[1]))
transition_matrix = self.evaluate(
tfd.Normal(0., 1.).sample([dim, dim], seed=model_seeds[2]))
transition_bias = self.evaluate(
tfd.Normal(0., 1.).sample(batch_shape + [dim],
seed=model_seeds[3]))
transition_scale_tril = self.evaluate(
tf.linalg.cholesky(
tfd.WishartTriL(
df=dim,
scale_tril=tf.eye(dim)).sample(seed=model_seeds[4])))
initial_state_prior = tfd.MultivariateNormalDiag(
loc=prior_loc, scale_diag=prior_scale, name='initial_state_prior')
lgssm = tfd.LinearGaussianStateSpaceModel(
num_timesteps=7,
transition_matrix=transition_matrix,
transition_noise=tfd.MultivariateNormalTriL(
loc=transition_bias, scale_tril=transition_scale_tril),
# Trivial observation model to pass through the latent state.
observation_matrix=tf.eye(dim),
observation_noise=tfd.MultivariateNormalDiag(
loc=tf.zeros(dim), scale_diag=tf.zeros(dim)),
initial_state_prior=initial_state_prior)
markov_chain = tfd.MarkovChain(
initial_state_prior=initial_state_prior,
transition_fn=lambda _, x: tfd.MultivariateNormalTriL( # pylint: disable=g-long-lambda
loc=tf.linalg.matvec(transition_matrix, x) + transition_bias,
scale_tril=transition_scale_tril),
num_steps=7)
x = markov_chain.sample(5, seed=seed)
self.assertAllClose(lgssm.log_prob(x),
markov_chain.log_prob(x),
rtol=1e-5)
def testCovariance(self):
mvn = tfd.MultivariateNormalDiag(
loc=tf.zeros([2, 3], dtype=tf.float32), validate_args=True)
self.assertAllClose(
self.evaluate(tf.eye(3, batch_shape=[2], dtype=tf.float32)),
self.evaluate(mvn.covariance()))
mvn = tfd.MultivariateNormalDiag(
loc=tf.zeros([3], dtype=tf.float32),
scale_identity_multiplier=[3., 2.],
validate_args=True)
self.assertAllEqual([2], mvn.batch_shape)
self.assertAllEqual([3], mvn.event_shape)
self.assertAllClose(
np.array([[[3., 0, 0], [0, 3, 0], [0, 0, 3]], [[2, 0, 0], [0, 2, 0],
[0, 0, 2]]])**2.,
self.evaluate(mvn.covariance()))
mvn = tfd.MultivariateNormalDiag(
loc=tf.zeros([3], dtype=tf.float32),
scale_diag=[[3., 2, 1], [4, 5, 6]],
validate_args=True)
self.assertAllEqual([2], mvn.batch_shape)
self.assertAllEqual([3], mvn.event_shape)
self.assertAllClose(
np.array([[[3., 0, 0], [0, 2, 0], [0, 0, 1]], [[4, 0, 0], [0, 5, 0],
[0, 0, 6]]])**2.,
self.evaluate(mvn.covariance()))
def testMultivariateNormalNd(self, event_size, num_samples):
def target_log_prob_fn(event):
return tfd.MultivariateNormalFullCovariance(
loc=tf.zeros(event_size),
covariance_matrix=tf.eye(event_size)).log_prob(event)
state = tf.zeros(event_size)
samples = []
for seed in range(num_samples):
[state], _, _ = no_u_turn_sampler.kernel(
target_log_prob_fn=target_log_prob_fn,
current_state=[state],
step_size=[0.3],
seed=seed)
npstate = state.numpy()
samples.append([npstate[0], npstate[1]])
samples = np.array(samples)
plt.scatter(samples[:, 0], samples[:, 1])
savefig(
self.model_dir, "projection_chain_{}d_normal_{}_steps.png".format(
event_size, num_samples))
plt.close()
target_samples = tfd.MultivariateNormalFullCovariance(
loc=tf.zeros(event_size),
covariance_matrix=tf.eye(event_size)).sample(num_samples,
seed=4).numpy()
plt.scatter(target_samples[:, 0], target_samples[:, 1])
savefig(
self.model_dir,
"projection_independent_{}d_normal_{}_samples.png".format(
event_size, num_samples))
plt.close()
def testMVNConjugateLinearUpdateSupportsBatchShape(self):
strm = test_util.test_seed_stream()
num_latents = 2
num_outputs = 4
batch_shape = [3, 1]
prior_mean = tf.ones([num_latents])
prior_scale = tf.eye(num_latents) * 5.
likelihood_scale = tf.linalg.LinearOperatorLowerTriangular(
tfb.ScaleTriL().forward(
tf.random.normal(shape=batch_shape +
[int(num_outputs * (num_outputs + 1) / 2)],
seed=strm())))
linear_transformation = tf.random.normal(
batch_shape + [num_outputs, num_latents], seed=strm()) * 5.
true_latent = tf.random.normal(batch_shape + [num_latents],
seed=strm())
observation = tf.linalg.matvec(linear_transformation, true_latent)
posterior_mean, posterior_prec = (tfd.mvn_conjugate_linear_update(
prior_mean=prior_mean,
prior_scale=prior_scale,
linear_transformation=linear_transformation,
likelihood_scale=likelihood_scale,
observation=observation))
self._mvn_linear_update_test_helper(
prior_mean=prior_mean,
prior_scale=prior_scale,
linear_transformation=linear_transformation,
likelihood_scale=likelihood_scale.to_dense(),
observation=observation,
candidate_posterior_mean=posterior_mean,
candidate_posterior_prec=posterior_prec.to_dense())
def testSoftToHardPermutation(self):
"""The solution of the matching for the identity matrix is range(N)."""
dims = 10
identity = tf.eye(dims)
result_matching = ed.layers.utils.soft_to_hard_permutation(identity)
result_matching_val = self.evaluate(result_matching)
self.assertAllEqual(result_matching_val[0], np.eye(dims))
def testInhomogeneousBackwards(self, scheme, accuracy_order):
# Tests solving du/dt = At + b for a backward time step.
# Compares with exact solution u(0) = exp(-At) u(t)
# + (exp(-At) - 1) A^(-1) b.
time_step = 0.0001
u = tf.constant([1, 2, -1, -2], dtype=tf.float64)
matrix = tf.constant(
[[1, -1, 0, 0], [3, 1, 2, 0], [0, -2, 1, 4], [0, 0, 3, 1]],
dtype=tf.float64)
b = tf.constant([1, -1, -2, 2], dtype=tf.float64)
tridiag_form = self._convert_to_tridiagonal_format(matrix)
actual = self.evaluate(
scheme(u, time_step, 0, lambda t: (tridiag_form, b)))
exponent = tf.linalg.expm(-matrix * time_step)
eye = tf.eye(4, 4, dtype=tf.float64)
u = tf.expand_dims(u, 1)
b = tf.expand_dims(b, 1)
expected = (
tf.matmul(exponent, u) +
tf.matmul(exponent - eye, tf.matmul(tf.linalg.inv(matrix), b)))
expected = self.evaluate(tf.squeeze(expected))
error_tolerance = 30 * time_step**(accuracy_order + 1)
self.assertLess(np.max(np.abs(actual - expected)), error_tolerance)
def test_fit_posterior_with_joint_q(self):
# Target distribution: equiv to MVNFullCovariance(cov=[[1., 1.], [1., 2.]])
def p_log_prob(z, x):
return tfd.Normal(0., 1.).log_prob(z) + tfd.Normal(z,
1.).log_prob(x)
# The Q family is a joint distribution that can express any 2D MVN.
b = tf.Variable([0., 0.])
l = tfp.util.TransformedVariable(tf.eye(2), tfb.ScaleTriL())
def trainable_q_fn():
z = yield tfd.JointDistributionCoroutine.Root(
tfd.Normal(b[0], l[0, 0], name='z'))
_ = yield tfd.Normal(b[1] + l[1, 0] * z, l[1, 1], name='x')
q = tfd.JointDistributionCoroutine(trainable_q_fn)
seed = tfp_test_util.test_seed()
loss_curve = tfp.vi.fit_surrogate_posterior(
p_log_prob,
q,
num_steps=1000,
sample_size=100,
optimizer=tf.optimizers.Adam(learning_rate=0.1),
seed=seed)
self.evaluate(tf1.global_variables_initializer())
loss_curve_ = self.evaluate((loss_curve))
# Since the Q family includes the true distribution, the optimized
# loss should be (approximately) zero.
self.assertAllClose(loss_curve_[-1], 0., atol=0.1)
def make_ar_transition_matrix(coefficients):
"""Build transition matrix for an autoregressive StateSpaceModel.
When applied to a vector of previous values, this matrix computes
the expected new value (summing the previous states according to the
autoregressive coefficients) in the top dimension of the state space,
and moves all previous values down by one dimension, 'forgetting' the
final (least recent) value. That is, it looks like this:
```
ar_matrix = [ coefs[0], coefs[1], ..., coefs[order]
1., 0 , ..., 0.
0., 1., ..., 0.
...
0., 0., ..., 1., 0. ]
```
Args:
coefficients: float `Tensor` of shape `concat([batch_shape, [order]])`.
Returns:
ar_matrix: float `Tensor` with shape `concat([batch_shape,
[order, order]])`.
"""
top_row = tf.expand_dims(coefficients, -2)
coef_shape = dist_util.prefer_static_shape(coefficients)
batch_shape, order = coef_shape[:-1], coef_shape[-1]
remaining_rows = tf.concat([
tf.eye(order - 1, dtype=coefficients.dtype, batch_shape=batch_shape),
tf.zeros(tf.concat([batch_shape, (order - 1, 1)], axis=0),
dtype=coefficients.dtype)
], axis=-1)
ar_matrix = tf.concat([top_row, remaining_rows], axis=-2)
return ar_matrix
def testDocstring1(self):
d = tfd.Blockwise(
[
tfd.Independent(tfd.Normal(loc=tf1.placeholder_with_default(
tf.zeros(4, dtype=tf.float64),
shape=None,
),
scale=1),
reinterpreted_batch_ndims=1),
tfd.MultivariateNormalTriL(
scale_tril=tf1.placeholder_with_default(
tf.eye(2, dtype=tf.float32), shape=None)),
],
dtype_override=tf.float32,
validate_args=True,
)
x = d.sample([2, 1], seed=42)
y = d.log_prob(x)
x_, y_ = self.evaluate([x, y])
self.assertEqual((2, 1, 4 + 2), x_.shape)
self.assertIs(tf.float32, x.dtype)
self.assertEqual((2, 1), y_.shape)
self.assertIs(tf.float32, y.dtype)
self.assertAllClose(np.zeros((6, ), dtype=np.float32),
self.evaluate(d.mean()))
def build_seasonal_transition_matrix(num_seasons,
is_last_day_of_season,
dtype,
basis_change_matrix=None,
basis_change_matrix_inv=None):
"""Build a function computing transitions for a seasonal effect model."""
with tf1.name_scope('build_seasonal_transition_matrix'):
# If the season is changing, the transition matrix permutes the latent
# state to shift all seasons up by a dimension, and sends the current
# season's effect to the bottom.
seasonal_permutation = np.concatenate([np.arange(1, num_seasons), [0]],
axis=0)
seasonal_permutation_matrix = tf.constant(
np.eye(num_seasons)[seasonal_permutation], dtype=dtype)
# Optionally transform the transition matrix into a reparameterized space,
# enforcing the zero-sum constraint for ConstrainedSeasonalStateSpaceModel.
if basis_change_matrix is not None:
seasonal_permutation_matrix = tf.matmul(
basis_change_matrix,
tf.matmul(seasonal_permutation_matrix,
basis_change_matrix_inv))
identity_matrix = tf.eye(
tf.shape(input=seasonal_permutation_matrix)[-1], dtype=dtype)
def seasonal_transition_matrix(t):
return tf.linalg.LinearOperatorFullMatrix(
matrix=dist_util.pick_scalar_condition(
is_last_day_of_season(t), seasonal_permutation_matrix,
identity_matrix))
return seasonal_transition_matrix
def test_log_prob_gradients_exist(self):
nsteps = 7
initial_mean = tf.constant([0.], dtype=self.dtype)
initial_cov = tf.eye(1, dtype=self.dtype)
transition_matrix = tf.constant(nsteps * [[[1.]]], dtype=self.dtype)
transition_mean = tf.constant(nsteps * [[0.]], dtype=self.dtype)
transition_cov = tf.constant(nsteps * [[[1.]]], dtype=self.dtype)
observation_matrix = tf.constant(nsteps * [[[1.]]], dtype=self.dtype)
observation_mean = tf.constant(nsteps * [[0.]], dtype=self.dtype)
observation_cov = tf.constant(nsteps * [[[1.]]], dtype=self.dtype)
y = tf.ones([nsteps, 1], dtype=self.dtype)
def lp(initial_mean, initial_cov, transition_matrix, transition_mean,
transition_cov, observation_matrix, observation_mean,
observation_cov, y):
(log_likelihoods, _, _, _, _, _,
_) = (parallel_kalman_filter_lib.kalman_filter(
transition_matrix=transition_matrix,
transition_mean=transition_mean,
transition_cov=transition_cov,
observation_matrix=observation_matrix,
observation_mean=observation_mean,
observation_cov=observation_cov,
initial_mean=initial_mean,
initial_cov=initial_cov,
y=y,
mask=None))
return tf.reduce_sum(log_likelihoods)
grads = tfp.math.value_and_gradient(
lp, (initial_mean, initial_cov, transition_matrix, transition_mean,
transition_cov, observation_matrix, observation_mean,
observation_cov, y))
for g in grads:
self.assertIsNotNone(g)
def eye(N, M=None, k=0, dtype=float): # pylint: disable=invalid-name,missing-docstring
if dtype:
dtype = utils.result_type(dtype)
if not M:
M = N
# Making sure N, M and k are `int`
N = int(N)
M = int(M)
k = int(k)
if k >= M or -k >= N:
# tf.linalg.diag will raise an error in this case
return zeros([N, M], dtype=dtype)
if k == 0:
return arrays_lib.tensor_to_ndarray(tf.eye(N, M, dtype=dtype))
# We need the precise length, otherwise tf.linalg.diag will raise an error
diag_len = min(N, M)
if k > 0:
if N >= M:
diag_len -= k
elif N + k > M:
diag_len = M - k
elif k <= 0:
if M >= N:
diag_len += k
elif M - k > N:
diag_len = N + k
diagonal = tf.ones([diag_len], dtype=dtype)
return arrays_lib.tensor_to_ndarray(
tf.linalg.diag(diagonal=diagonal, num_rows=N, num_cols=M, k=k))
def newton_qr(jacobian, newton_coefficient, step_size): """QR factorizes the matrix used in each iteration of Newton's method.""" identity = tf.eye(tf.shape(jacobian)[0], dtype=jacobian.dtype) step_size_cast = tf.cast(step_size, jacobian.dtype) newton_matrix = (identity - step_size_cast * newton_coefficient * jacobian) factorization = tf.linalg.qr(newton_matrix) return factorization.q, factorization.r
def assert_mvn_target_conservation(event_size, batch_size, **kwargs):
strm = tfp_test_util.test_seed_stream()
initialization = tfd.MultivariateNormalFullCovariance(
loc=tf.zeros(event_size),
covariance_matrix=tf.eye(event_size)).sample(
batch_size, seed=strm())
samples, _ = run_nuts_chain(
event_size, batch_size, num_steps=1,
initial_state=initialization, **kwargs)
answer = samples[0][-1]
check_cdf_agrees = (
st.assert_multivariate_true_cdf_equal_on_projections_two_sample(
answer, initialization, num_projections=100, false_fail_rate=1e-6))
check_sample_shape = assert_util.assert_equal(
tf.shape(answer)[0], batch_size)
movement = tf.linalg.norm(answer - initialization, axis=-1)
# This movement distance (0.3) was copied from the univariate case.
check_movement = assert_util.assert_greater_equal(
tf.reduce_mean(movement), 0.3)
check_enough_power = assert_util.assert_less(
st.min_discrepancy_of_true_cdfs_detectable_by_dkwm_two_sample(
batch_size, batch_size, false_fail_rate=1e-8, false_pass_rate=1e-6),
0.055)
return (
check_cdf_agrees,
check_sample_shape,
check_movement,
check_enough_power,
)
def create_hmm(self, num_steps):
"""Same as the original CREPE viterbi decdoding, but in TF."""
# Initial distribution is uniform.
initial_distribution = tfp.distributions.Categorical(
probs=tf.ones([360]) / 360)
# Transition probabilities inducing continuous pitch.
bins = tf.range(360, dtype=tf.float32)
xx, yy = tf.meshgrid(bins, bins)
min_transition = 1e-5 # For training stabiity.
transition = tf.maximum(12 - abs(xx - yy), min_transition)
transition = transition / tf.reduce_sum(transition, axis=1)[:, None]
transition = tf.cast(transition, tf.float32)
transition_distribution = tfp.distributions.Categorical(
probs=transition)
# Emission probability = fixed probability for self, evenly distribute the
# others.
self_emission = 0.1
emission = (tf.eye(360) * self_emission + tf.ones(shape=(360, 360)) *
((1 - self_emission) / 360.))
emission = tf.cast(emission, tf.float32)[None, ...]
observation_distribution = tfp.distributions.Multinomial(
total_count=1, probs=emission)
return tfp.distributions.HiddenMarkovModel(
initial_distribution=initial_distribution,
transition_distribution=transition_distribution,
observation_distribution=observation_distribution,
num_steps=num_steps,
)
def call(self, inputs, logits=None, training=None):
"""Minibatch updates the GP's posterior precision matrix estimate.
Args:
inputs: (tf.Tensor) GP random features, shape (batch_size,
gp_hidden_size).
logits: (tf.Tensor) Pre-activation output from the model. Needed
for Laplace approximation under a non-Gaussian likelihood.
training: (tf.bool) whether or not the layer is in training mode. If in
training mode, the gp_weight covariance is updated using gp_feature.
Returns:
gp_stddev (tf.Tensor): GP posterior predictive variance,
shape (batch_size, batch_size).
"""
batch_size = tf.shape(inputs)[0]
training = self._get_training_value(training)
if training:
# Define and register the update op for feature precision matrix.
precision_matrix_update_op = self.make_precision_matrix_update_op(
gp_feature=inputs,
logits=logits,
precision_matrix=self.precision_matrix)
self.add_update(precision_matrix_update_op)
# Return null estimate during training.
return tf.eye(batch_size, dtype=self.dtype)
else:
# Return covariance estimate during inference.
return self.compute_predictive_covariance(gp_feature=inputs)
def step(per_replica_inputs: _TensorDict) -> None:
"""The function defining a single validation/test step."""
features = per_replica_inputs['features']
labels = per_replica_inputs[label_key]
logits = model(features, training=False)
if isinstance(logits, (tuple, list)):
logits, covmat = logits
else:
per_core_batch_size, _ = logits.get_shape().as_list()
covmat = tf.eye(per_core_batch_size)
logits = ed.layers.utils.mean_field_logits(
logits, covmat, mean_field_factor=mean_field_factor)
predictions = tf.nn.softmax(logits, axis=-1)
if label_key != 'labels':
predictions = tf.reduce_max(predictions, axis=-1)
# Later when metric.result() is called, it will return the computed
# result, averaged across replicas.
for metric in metrics.values():
if isinstance(metric, tf.keras.metrics.Metric):
metric.update_state(labels, predictions) # pytype: disable=attribute-error
else:
metric.add_batch(predictions, label=labels)
return
def _forward(self, x):
with tf.control_dependencies(self._assertions(x)):
shape = tf.shape(x)
return tf.linalg.triangular_solve(
x,
tf.eye(shape[-1], batch_shape=shape[:-2], dtype=x.dtype),
lower=True)
def testCovariance(self):
vex = tfd.VectorExponentialDiag(loc=tf.ones([2, 3], dtype=tf.float32),
validate_args=True)
self.assertAllClose(
self.evaluate(tf.eye(3, batch_shape=[2], dtype=np.float32)),
self.evaluate(vex.covariance()))
vex = tfd.VectorExponentialDiag(loc=tf.ones([3], dtype=tf.float32),
scale_identity_multiplier=[3., 2.],
validate_args=True)
self.assertAllEqual([2], vex.batch_shape)
self.assertAllEqual([3], vex.event_shape)
self.assertAllClose(
np.array([[[3., 0, 0], [0, 3, 0], [0, 0, 3]],
[[2, 0, 0], [0, 2, 0], [0, 0, 2]]])**2.,
self.evaluate(vex.covariance()))
vex = tfd.VectorExponentialDiag(loc=tf.ones([3], dtype=tf.float32),
scale_diag=[[3., 2, 1], [4, 5, 6]],
validate_args=True)
self.assertAllEqual([2], vex.batch_shape)
self.assertAllEqual([3], vex.event_shape)
self.assertAllClose(
np.array([[[3., 0, 0], [0, 2, 0], [0, 0, 1]],
[[4, 0, 0], [0, 5, 0], [0, 0, 6]]])**2.,
self.evaluate(vex.covariance()))
def assert_mvn_target_conservation(event_size, batch_size, **kwargs):
initialization = tfd.MultivariateNormalFullCovariance(
loc=tf.zeros(event_size),
covariance_matrix=tf.eye(event_size)).sample(batch_size, seed=4)
samples, leapfrogs = run_nuts_chain(event_size,
batch_size,
num_steps=1,
initial_state=initialization,
**kwargs)
answer = samples[0][-1]
check_cdf_agrees = (
st.assert_multivariate_true_cdf_equal_on_projections_two_sample(
answer, initialization, num_projections=100, false_fail_rate=1e-6))
check_sample_shape = tf1.assert_equal(
tf.shape(input=answer)[0], batch_size)
unique, _ = tf.unique(leapfrogs[0])
check_leapfrogs_vary = tf1.assert_greater_equal(
tf.shape(input=unique)[0], 3)
avg_leapfrogs = tf.math.reduce_mean(input_tensor=leapfrogs[0])
check_leapfrogs = tf1.assert_greater_equal(
avg_leapfrogs, tf.constant(4, dtype=avg_leapfrogs.dtype))
movement = tf.linalg.norm(tensor=answer - initialization, axis=-1)
# This movement distance (0.3) was copied from the univariate case.
check_movement = tf1.assert_greater_equal(
tf.reduce_mean(input_tensor=movement), 0.3)
check_enough_power = tf1.assert_less(
st.min_discrepancy_of_true_cdfs_detectable_by_dkwm_two_sample(
batch_size, batch_size, false_fail_rate=1e-8,
false_pass_rate=1e-6), 0.055)
return (check_cdf_agrees, check_sample_shape, check_leapfrogs_vary,
check_leapfrogs, check_movement, check_enough_power)
def testEigvalsAsserts(self):
with self.assertRaisesOpError(r'must be positive'):
dpp = tfd.DeterminantalPointProcess(tf.constant(
[1, 2, 3, 0.], dtype=self.param_dtype),
tf.eye(4,
dtype=self.param_dtype),
validate_args=True)
self.evaluate(dpp.sample(seed=test_util.test_seed()))
v = tf.Variable(tf.constant([1, 2, -3, 4.], dtype=self.param_dtype))
self.evaluate(v.initializer)
dpp = tfd.DeterminantalPointProcess(v,
tf.eye(4, dtype=self.param_dtype),
validate_args=True)
with self.assertRaisesOpError(r'must be positive'):
self.evaluate(dpp.sample(seed=test_util.test_seed()))
def _embed_l2_regularization():
"""Adds synthetic observations to implement L2 regularization."""
# `tf.matrix_solve_ls` does not respect the `l2_regularization` argument
# when `fast_unsafe_numerics` is `False`. This function adds synthetic
# observations to the data to implement the regularization instead.
# Adding observations `sqrt(l2_regularizer) * I` is mathematically
# equivalent to adding the term
# `-l2_regularizer ||coefficients||_2**2` to the log-likelihood.
num_model_coefficients = num_cols(model_matrix)
batch_shape = tf.shape(model_matrix)[:-2]
if l2_regularization_penalty_factor is None:
eye = tf.eye(
num_model_coefficients, batch_shape=batch_shape, dtype=a.dtype)
else:
eye = tf.linalg.tensor_diag(
tf.cast(l2_regularization_penalty_factor, dtype=a.dtype))
broadcasted_shape = prefer_static.concat(
[batch_shape, [num_model_coefficients, num_model_coefficients]],
axis=0)
eye = tf.broadcast_to(eye, broadcasted_shape)
a_ = tf.concat([a, tf.sqrt(l2_regularizer) * eye], axis=-2)
b_ = distribution_util.pad(
b, count=num_model_coefficients, axis=-1, back=True)
# Return l2_regularizer=0 since its now embedded.
l2_regularizer_ = np.array(0, dtype_util.as_numpy_dtype(a.dtype))
return a_, b_, l2_regularizer_
def test_weighted_sum_tensor(self, dtype):
del dtype # not used in this test case.
weights = [0.5, -0.25, -0.25]
states = [tf.eye(2) for _ in range(3)]
weighted_tensor_sum = rk_util.weighted_sum(weights, states)
self.assertAllClose(weighted_tensor_sum, tf.zeros((2, 2)))
weights = [0.5, -0.25, -0.25, 1.0]
states = [tf.ones(2) for _ in range(4)]
weighted_tensor_sum = rk_util.weighted_sum(weights, states)
self.assertAllClose(weighted_tensor_sum, tf.ones(2))
weights = [0.5, -0.25, -0.25, 0.0]
states = [tf.eye(2) for _ in range(4)]
weighted_tensor_sum = rk_util.weighted_sum(weights, states)
self.assertAllClose(weighted_tensor_sum, tf.zeros((2, 2)))
