def test_discrete_mixture_forward_random(self):
"""Tests if forward pass with a randomly generated mask matches outputs from the selected flows."""
N, K, B = self.N, self.K, self.B
base_samples = self.base_samples
for i, specs in enumerate(self.mixture_specs):
logger.debug("#"*30+" [%i/%i] testing forward of mixture of %s" % \
(i+1, len(self.mixture_specs), specs))
flows1 = flows.parse_layers_specification(specs, N, K, temperature=5.0)
mixture = flows_mixture.DiscreteFlowsMixture(N, K, B, flows=flows1, temperature=5.0) #mixture of the flows
bs = np.random.choice(range(B), base_samples.shape[0]) # assign each sample to a flow from mixture
mask = tf.one_hot(bs, B)[...,None,:,None] # create mask
out_samples = mixture(base_samples, mask=mask)
self.assertTrue( bool(tf.reduce_all(tf.reduce_sum(out_samples, -1)==1)),
"masked samples passed through mixture=%s have exactly one 1 per dim." % (specs))
for b in range(B):
bth_flow_sample_indices = np.nonzero(bs==b)[0]
bth_flow_out_samples = tf.gather(out_samples, bth_flow_sample_indices, -2)
# pass assigned samples through individual flow
flow1_out_samples = flows1[b](base_samples[...,b,:])
bth_flow1_out_samples = tf.gather(flow1_out_samples, bth_flow_sample_indices, -2)
self.assertTrue(bth_flow_out_samples.shape==bth_flow1_out_samples.shape,
"mask selects correct number of samples assigned to %d-th flow" % b)
self.assertTrue(bool(tf.reduce_all(bth_flow_out_samples==bth_flow1_out_samples)),
"check if results from a mixture match results from individual %d-th flow" % b)
def test_discrete_mixture_forward_nonmasked(self):
"""Tests if forward pass without masking results in invalid samples having exactly B ones per dimension."""
N, K, B = self.N, self.K, self.B
base_samples = self.base_samples
for i, specs in enumerate(self.mixture_specs):
logger.debug("#"*30+" [%i/%i] testing forward (no mask) of mixture of %s" % \
(i+1, len(self.mixture_specs), specs))
flows1 = flows.parse_layers_specification(specs, N, K, temperature=5.0)
mixture = flows_mixture.DiscreteFlowsMixture(N, K, B, flows=flows1, temperature=5.0) #mixture of the flows
out_samples = mixture(base_samples)
self.assertTrue( bool(tf.reduce_all(tf.reduce_sum(out_samples, -1)==B)),
"non-masked transformation through mixture=%s has %i ones per dim." % (specs, B))
def __init__(self,
N,
K,
B,
flows=None,
temperature=None,
single_flow_layers=[("M", [64, 64])],
components_specification=None,
**kwargs):
"""
Args:
flows A list of components
temperature ST hyperparameter
single_flow_layers Specification of layers.
If used all components will follow the same design.
Either a list of flows or a temperature for newly created
(according to single_flow_layers) flows should be provided.
"""
super(DiscreteFlowsMixture, self).__init__(**kwargs)
self._N = N
self._K = K
self._B = B
assert temperature is not None or flows is not None, \
"You have to provide either a list of flows or a temperature for newly created flows!"
if flows is None:
if components_specification is not None:
assert len(components_specification)==B, \
"components_specification should contain one layer specification per flow"
flows = parse_layers_specification(components_specification,
N,
K,
temperature,
dtype=self.dtype)
elif single_flow_layers is not None:
flows = [
DiscreteFlow(N, K, temperature, layers=single_flow_layers)
for _ in range(B)
]
else:
raise ValueError("You have to fix one of the following:" + \
"flows, components_specification or single_flow_layers!")
assert len(flows) == B
self.flows = flows
def test_discrete_mixture_reverse(self):
"""Tests if samples passed in the reverse direction consist of reversed samples from individual flows."""
N, K, B = self.N, self.K, self.B
base_samples = self.base_samples
for i, specs in enumerate(self.mixture_specs):
logger.debug("#"*30+" [%i/%i] testing reverse of mixture of %s" % \
(i+1, len(self.mixture_specs), specs))
flows1 = flows.parse_layers_specification(specs, N, K, temperature=5.0)
mixture = flows_mixture.DiscreteFlowsMixture(N,K,B,flows=flows1, temperature=5.0) #mixture of the flows
# test if reverse is composed out of reversed of individual
out_samples = base_samples[...,0,:] # just any 'output' x
inv_samples = mixture.reverse(out_samples) # reverse transformation to u
for b, f in enumerate(flows1):
self.assertTrue( bool(tf.reduce_all(f.reverse(out_samples)==inv_samples[...,b,:])), \
"reverse of %i-th flow" % b)
def test_discrete_mixture_forward(self):
"""Tests if forward pass with a mask matches outputs from the mask-selected flows."""
N, K, B = self.N, self.K, self.B
base_samples = self.base_samples
for i, specs in enumerate(self.mixture_specs):
logger.debug("#"*30+" [%i/%i] testing forward of mixture of %s" % \
(i+1, len(self.mixture_specs), specs))
flows1 = flows.parse_layers_specification(specs, N, K, temperature=5.0)
mixture = flows_mixture.DiscreteFlowsMixture(N, K, B, flows=flows1, temperature=5.0) #mixture of the flows
for b in range(B):
flow1_out_samples = flows1[b](base_samples[...,b,:])
mask = tf.one_hot([b], B)[None,:][...,None]
out_samples = mixture(base_samples, mask=mask)
self.assertTrue( bool(tf.reduce_all(tf.reduce_sum(out_samples, -1)==1)),
"masked samples passed through mixture=%s have exactly one 1 per dim." % (specs))
self.assertTrue( tf.reduce_all(flow1_out_samples == out_samples),
"choosing outputs from %i-th flow" % b)