def get_samples(distribution, num_samples, seed=None):
  """Given a batched distribution, compute samples and reshape along batch.

  That is, we have a distribution of shape (batch_size, ...), where each element
  of the tensor is independent. We then draw num_samples from each component, to
  give a tensor of shape:

      (num_samples, batch_size, ...)

  Args:
    distribution: `tfp.distributions.Distribution`. The distribution from which
      to sample.
    num_samples: `Integral` | `DeviceArray`, int32, (). The number of samples.
    seed: `Integral` | `None`. The seed that will be forwarded to the call to
      distribution.sample. Defaults to `None`.

  Returns:
    `DeviceArray`, float32, (batch_size * num_samples, ...).
    Samples for each element of the batch.
  """
  # Obtain the sample from the distribution, which will be of shape
  # [num_samples] + batch_shape + event_shape.
  sample = distribution.sample(num_samples, seed=seed)
  sample = sample.reshape((-1, sample.shape[-1]))

  # Combine the first two dimensions through a reshape, so the result will
  # be of shape (num_samples * batch_size,) + shape_tail.
  return sample