def new(params, event_size, dtype=None, validate_args=False, name=None):
"""Create the distribution instance from a `params` vector."""
with tf.name_scope(name, 'OneHotCategorical', [params, event_size]):
return tfd.OneHotCategorical(logits=params,
dtype=dtype
or params.dtype.base_dtype,
validate_args=validate_args)
def new(params,
event_size,
num_components,
dtype=None,
validate_args=False,
name=None):
"""Create the distribution instance from a `params` vector."""
with tf.name_scope(name, 'CategoricalMixtureOfOneHotCategorical',
[params, event_size, num_components]):
components_shape = tf.concat(
[tf.shape(params)[:-1], [num_components, event_size]], axis=0)
dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(
logits=params[..., :num_components],
validate_args=validate_args),
components_distribution=tfd.OneHotCategorical(
logits=tf.reshape(params[..., num_components:],
components_shape),
dtype=dtype or params.dtype.base_dtype,
validate_args=False
), # So we can eval on simplex interior.
# TODO(b/120154797): Change following to `validate_args=True` after
# fixing: "ValueError: `mixture_distribution` must have scalar
# `event_dim`s." assertion in MixtureSameFamily.
validate_args=False)
# pylint: disable=protected-access
dist._mean = functools.partial(_eval_all_one_hot,
tfd.Distribution.prob, dist)
dist.log_mean = functools.partial(_eval_all_one_hot,
tfd.Distribution.log_prob, dist)
# pylint: enable=protected-access
return dist
def f_child_sample(targets: List[Tuple[tf.Tensor, ts.NestedTensor]]) -> ts.NestedTensor:
# select bias for top_down attention
dist = tfd.OneHotCategorical(logits=[beta-energy for energy, latent in targets])
sample = dist.sample()
ind = tf.argmax(sample)
energy, latent = targets[ind]
energy = energy + dist.entropy() + -dist.log_prob(sample)
return energy, latent
def f_parent_sample(states: ts.NestedTensor, parent_names: List[Text]):
# select parent to pool data from by energy weighted bottom-up attention
parent_energies = [states[name][keys.STATES.ENERGY] for name in parent_names]
parent_dist = tfd.OneHotCategorical(logits=parent_energies)
parent_name = parent_names[tf.argmax(parent_dist.sample())]
parent_energy = states[parent_name][keys.STATES.ENERGY]
parent_latent = states[parent_name][keys.STATES.LATENT]
return parent_energy, parent_latent
def new(params,
probs_input=False,
dtype=None,
validate_args=False,
name=None):
"""Create the distribution instance from a `params` vector."""
with tf.compat.v1.name_scope(name, 'OneHotCategorical', [params]):
return tfd.OneHotCategorical(
logits=params if not probs_input else None,
probs=tf.clip_by_value(params, 1e-8, 1 -
1e-8) if probs_input else None,
dtype=dtype or params.dtype.base_dtype,
validate_args=validate_args)
def new(params,
probs_input=False,
dtype=None,
validate_args=False,
name='OneHotCategoricalLayer'):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
return tfd.OneHotCategorical(
logits=params if not probs_input else None,
probs=tf.clip_by_value(params, 1e-8, 1 - 1e-8) \
if probs_input else None,
dtype=dtype or params.dtype,
validate_args=validate_args,
name=name)
def new(params,
probs_input=False,
dtype=None,
validate_args=False,
name='OneHotCategoricalLayer'):
r"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
if probs_input:
logits = None
probs = params
else:
logits = params
probs = None
dist = tfd.OneHotCategorical(
logits=logits,
probs=probs,
dtype=dtype or params.dtype,
validate_args=validate_args,
name=name)
return dist
def get_distributions(self, validate_args=False):
self.dist1 = tfd.MultivariateNormalDiag(loc=self.maybe_static(
tf.zeros(self.batch_dim_1 + self.event_dim_1, dtype=self.dtype),
self.is_static),
scale_diag=self.maybe_static(
tf.ones(self.batch_dim_1 +
self.event_dim_1,
dtype=self.dtype),
self.is_static))
self.dist2 = tfd.OneHotCategorical(logits=self.maybe_static(
tf.zeros(self.batch_dim_2 + self.event_dim_2), self.is_static),
dtype=self.dtype)
self.dist3 = tfd.Dirichlet(
self.maybe_static(
tf.zeros(self.batch_dim_3 + self.event_dim_3,
dtype=self.dtype), self.is_static))
return batch_concat.BatchConcat(
distributions=[self.dist1, self.dist2, self.dist3],
axis=self.axis,
validate_args=validate_args)