def __call__(self, shape, dtype=None):
if not self.built:
self.build(shape, dtype)
loc = self.loc
if self.loc_constraint:
loc = self.loc_constraint(loc)
return generated_random_variables.Independent(
generated_random_variables.Deterministic(loc=loc).distribution,
reinterpreted_batch_ndims=len(shape))
def __call__(self, shape, dtype=None):
if not self.built:
self.build(shape, dtype)
loc = self.loc
if self.loc_constraint:
loc = self.loc_constraint(loc)
return generated_random_variables.Independent(
generated_random_variables.MixtureSameFamily(
mixture_distribution=generated_random_variables.Categorical(
probs=tf.broadcast_to(
[[1/self.num_components]*self.num_components],
list(shape) + [self.num_components])).distribution,
components_distribution=generated_random_variables.Deterministic(
loc=loc).distribution
).distribution,
reinterpreted_batch_ndims=len(shape))