code: def batched_multinomial(
      weights: Tensor,
      num_samples: int,
      replacement: bool = False,
      generator: Optional[torch.Generator] = None,
      out: Optional[Tensor] = None,
  ) -> LongTensor:
      r"""Sample from multinomial with an arbitrary number of batch dimensions.
  
      Args:
          weights: A `batch_shape x num_categories` tensor of weights. For each batch
              index `i, j, ...`, this functions samples from a multinomial with `input`
              `weights[i, j, ..., :]`. Note that the weights need not sum to one, but must
              be non-negative, finite and have a non-zero sum.
          num_samples: The number of samples to draw for each batch index. Must be smaller
              than `num_categories` if `replacement=False`.
          replacement: If True, samples are drawn with replacement.
          generator: A a pseudorandom number generator for sampling.
          out: The output tensor (optional). If provided, must be of size
              `batch_shape x num_samples`.
  
      Returns:
          A `batch_shape x num_samples` tensor of samples.
  
      This is a thin wrapper around `torch.multinomial` that allows weight (`input`)
      tensors with an arbitrary number of batch dimensions (`torch.multinomial` only
      allows a single batch dimension). The calling signature is the same as for
      `torch.multinomial`.
  
      Example:
          >>> weights = torch.rand(2, 3, 10)
          >>> samples = batched_multinomial(weights, 4)  # shape is 2 x 3 x 4
      """
      batch_shape, n_categories = weights.shape[:-1], weights.size(-1)
      flat_samples = torch.multinomial(
          input=weights.view(-1, n_categories),
          num_samples=num_samples,
          replacement=replacement,
          generator=generator,
          out=None if out is None else out.view(-1, num_samples),
      )
      return flat_samples.view(*batch_shape, num_samples)
  
imports: from __future__ import annotations
  from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from typing import Any, Generator, Iterable, List, Optional, Tuple, TYPE_CHECKING
  import numpy
  import scipy
  import torch
  from botorch.exceptions.errors import BotorchError
  from botorch.sampling.qmc import NormalQMCEngine
  from botorch.utils.transforms import unnormalize
  from scipy.spatial import Delaunay, HalfspaceIntersection
  from torch import LongTensor, Tensor
  from torch.quasirandom import SobolEngine
label_desc: Use image node_id 7 for calling a global function with example usage: batched_multinomial(weights, num_samples, replacement, generator, out) and returns: flat_samples