'Input is split into a list of chunks according
to dist.param_dim along axis=1
Args:
s (tf.Tensor): batch of vectors with shape (batch, param_dim or sample_dim)
param (bool): split params, otherwise split samples
Yields:
tf.Tensor: chunk from input of length N_i with sum N_i = N'

def _splitter(self, s, param):

    offset = 0
    for dist in self.dists:
        if param:
            off = dist.param_dim
        else:
            off = dist.sample_dim
        (yield s[:, offset:(offset + off)])
        offset += off
