def _probs_for_boltzmann_layer(group_size, p):
# Shortcut for degenerate case.
if group_size == 1:
return sigmoid(p, out=p)
p.shape = p.shape[:-1] + (-1, group_size)
boltzmann_dist(p, axis=-1, out=p)
p.shape = p.shape[:-2] + (-1,)
return p
def _sample_boltzmann_layer(group_size, s):
# Shortcut for degenerate case.
if group_size == 1:
return sample_indicator(sigmoid(s, out=s), out=s)
s.shape = s.shape[:-1] + (-1, group_size)
boltzmann_dist(s, axis=-1, out=s)
sample_exclusive_indicators(s, axis=-1, out=s)
s.shape = s.shape[:-2] + (-1,)
return s
