code: def sample_pmf(pmf: torch.Tensor, n_samples: int) -> torch.Tensor:
      r"""
      Sample from the given discrete probability distribution with replacement.
  
      The i-th bin is assumed to have mass pmf[i].
  
      Args:
          pmf: [batch_size, *shape, n_samples, 1] where (pmf.sum(dim=-2) == 1).all()
          n_samples: number of samples
  
      Return:
          indices sampled with replacement
      """
  
      *shape, support_size, last_dim = pmf.shape
      assert last_dim == 1
  
      cdf = torch.cumsum(pmf.view(-1, support_size), dim=1)
      inds = torch.searchsorted(
          cdf, torch.rand(cdf.shape[0], n_samples, device=cdf.device)
      )
  
      return inds.view(*shape, n_samples, 1).clamp(0, support_size - 1)
  
imports: import math
  from dataclasses import dataclass
  from typing import Dict, Optional, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from torch import nn
  from ...configuration_utils import ConfigMixin, register_to_config
  from ...models import ModelMixin
  from ...utils import BaseOutput
  from .camera import create_pan_cameras
label_desc: Use image node_id 1 for calling a global function with example usage: sample_pmf(pmf, n_samples) and returns: inds