docs: def sample_discrete(values, weights, size=1):
      """
      Generate a sample of given size, given a probability mass function.
  
      Parameters
      ----------
      values: numpy.array: Array of all possible values that the random variable
              can take.
      weights: numpy.array or list of numpy.array: Array(s) representing the PMF of the random variable.
      size: int: Size of the sample to be generated.
  
      Returns
      -------
      numpy.array: of values of the random variable sampled from the given PMF.
  
      Example
      -------
      >>> import numpy as np
      >>> from pgmpy.utils.mathext import sample_discrete
      >>> values = np.array(['v_0', 'v_1', 'v_2'])
      >>> probabilities = np.array([0.2, 0.5, 0.3])
      >>> sample_discrete(values, probabilities, 10)
      array(['v_1', 'v_1', 'v_0', 'v_1', 'v_2', 'v_0', 'v_1', 'v_1', 'v_1',
        'v_2'], dtype='<U3')
      """
      weights = np.array(weights)
      if weights.ndim == 1:
          return np.random.choice(values, size=size, p=weights)
      else:
          return np.fromiter(map(lambda t: np.random.choice(values, p=t), weights), dtype='int')