'Sample a batch of vectors from this distribution parameterized by theta.
Args:
batch_size(int): the batch size.
theta: a tensor of shape (param_dim,) or (batch, param_dim).
Returns:
a batch of samples of shape (batch, sample_dim)'

@class_scope
def sample(self, batch_size, theta):

    assert isinstance(batch_size, int), batch_size
    shp = theta.get_shape()
    assert ((shp.ndims in [1, 2]) and (shp[(-1)] == self.sample_dim)), shp
    if (shp.ndims == 1):
        theta = tf.tile(tf.expand_dims(theta, 0), [batch_size, 1], name='tiled_theta')
    else:
        assert (shp[0] == batch_size), shp
    x = self._sample(batch_size, theta)
    assert ((x.get_shape().ndims == 2) and (x.get_shape()[1] == self.sample_dim)), x.get_shape()
    return x
