code: def _safe_multinomial(
      input: torch.Tensor, num_samples: int
  ) -> torch.Tensor:
      """
      Wrapper around torch.multinomial that attempts sampling without replacement
      when possible, otherwise resorts to sampling with replacement.
  
      Args:
          input: tensor of shape [B, n] containing non-negative values;
                  rows are interpreted as unnormalized event probabilities
                  in categorical distributions.
          num_samples: number of samples to take.
  
      Returns:
          LongTensor of shape [B, num_samples] containing
          values from {0, ..., n - 1} where the elements [i, :] of row i make
              (1) if there are num_samples or more non-zero values in input[i],
                  a random subset of the indices of those values, with
                  probabilities proportional to the values in input[i, :].
  
              (2) if not, a random sample with replacement of the indices of
                  those values, with probabilities proportional to them.
                  This sample might not contain all the indices of the
                  non-zero values.
          Behavior undetermined if there are no non-zero values in a whole row
          or if there are negative values.
      """
      try:
          res = torch.multinomial(input, num_samples, replacement=False)
      except RuntimeError:
          # this is probably rare, so we don't mind sampling twice
          res = torch.multinomial(input, num_samples, replacement=True)
          no_repl = (input > 0.0).sum(dim=-1) >= num_samples
          res[no_repl] = torch.multinomial(
              input[no_repl], num_samples, replacement=False
          )
          return res
  
      # in some versions of Pytorch, zero probabilty samples can be drawn without an error
      # due to this bug: https://github.com/pytorch/pytorch/issues/50034. Handle this case:
      repl = (input > 0.0).sum(dim=-1) < num_samples
      if repl.any():
          res[repl] = torch.multinomial(
              input[repl], num_samples, replacement=True
          )
  
      return res
  
imports: import warnings
  from typing import Optional, Tuple, Union
  import torch
  from pytorch3d.common.compat import meshgrid_ij
  from pytorch3d.ops import padded_to_packed
  from pytorch3d.renderer.cameras import CamerasBase
  from pytorch3d.renderer.implicit.utils import HeterogeneousRayBundle, RayBundle
  from torch.nn import functional
label_desc: Use image node_id 3 for calling a global function with example usage: _safe_multinomial(input, num_samples) and returns: res, res