code: def sample_multinomial(N, p, size=None):
      r
      # ensure s is array
      s = np.array([1]) if size is None else np.array([size]).flatten()
  
      def take_samples(ps):
          # we have to flatten to make apply_along_axis work.
          return np.random.multinomial(N, ps, np.prod(s)).flatten()
  
      # should have shape (prod(size)*ps.shape[0], ps.shape[1:])
      samples = np.apply_along_axis(take_samples, 0, p)
      # should have shape (size, p.shape)
      samples = samples.reshape(np.concatenate([s, p.shape]))
      # should have shape (p.shape, size)
      samples = samples.transpose(np.concatenate(
          [np.arange(s.ndim, p.ndim+s.ndim), np.arange(s.ndim)]
      ))
  
      if size is None:
          # get rid of trailing singleton dimension.
          samples = samples[...,0]
  
      return samples
docs: r"""
      Draws fixed number of samples N from different
      multinomial distributions (with the same number dice sides).
  
      :param int N: How many samples to draw from each distribution.
      :param np.ndarray p: Probabilities specifying each distribution.
          Sum along axis 0 should be 1.
      :param size: Output shape. ``int`` or tuple of
          ``int``s. If the given shape is,
          e.g., ``(m, n, k)``, then m * n * k samples are drawn
          for each distribution.
          Default is None, in which case a single value
          is returned for each distribution.
  
      :rtype: np.ndarray
      :return: Array of shape ``(p.shape, size)`` or p.shape if
          size is ``None``.