code: def sample(
      self, rng: np.random.RandomState, batch_size: int
  ) -> np.ndarray:
      active_indices = rng.choice(
          len(self.weights), p=self.weights, size=batch_size
      )
  
      ret = np.empty(
          (batch_size, len(self.distributions)), dtype=np.float64
      )
      for i, d in enumerate(self.distributions):
          if isinstance(d, _BatchedCategoricalDistributions):
              active_weights = d.weights[active_indices, :]
              rnd_quantile = rng.rand(batch_size)
              cum_probs = np.cumsum(active_weights, axis=-1)
              assert np.isclose(cum_probs[:, -1], 1).all()
              cum_probs[:, -1] = 1  # Avoid numerical errors.
              ret[:, i] = np.sum(
                  cum_probs < rnd_quantile[:, None], axis=-1
              )
          elif isinstance(d, _BatchedTruncNormDistributions):
              active_mus = d.mu[active_indices]
              active_sigmas = d.sigma[active_indices]
              ret[:, i] = _truncnorm.rvs(
                  a=(d.low - active_mus) / active_sigmas,
                  b=(d.high - active_mus) / active_sigmas,
                  loc=active_mus,
                  scale=active_sigmas,
                  random_state=rng,
              )
          elif isinstance(d, _BatchedDiscreteTruncNormDistributions):
              active_mus = d.mu[active_indices]
              active_sigmas = d.sigma[active_indices]
              samples = _truncnorm.rvs(
                  a=(d.low - d.step / 2 - active_mus) / active_sigmas,
                  b=(d.high + d.step / 2 - active_mus) / active_sigmas,
                  loc=active_mus,
                  scale=active_sigmas,
                  random_state=rng,
              )
              ret[:, i] = np.clip(
                  d.low + np.round((samples - d.low) / d.step) * d.step,
                  d.low,
                  d.high,
              )
          else:
              assert False
  
      return ret
  
file_path: optuna/optuna/samplers/_tpe/probability_distributions.py
imports: from typing import List
  from typing import NamedTuple
  from typing import Union
  import numpy
  from optuna.samplers._tpe import _truncnorm
label_desc: Use image node_id 1 for calling the _MixtureOfProductDistribution obj's underlying member method code with example usage: obj.sample(rng, batch_size) and returns: ret