class_docstr: Samples a fixed number of points along rays which are regularly distributed
  in a batch of rectangular image grids. Points along each ray
  have uniformly-spaced z-coordinates between a predefined
  minimum and maximum depth.
  
  The raysampler first generates a 3D coordinate grid of the following form::
  
         / min_x, min_y, max_depth -------------- / max_x, min_y, max_depth
        /                                        /|
       /                                        / |     ^
      / min_depth                    min_depth /  |     |
      min_x ----------------------------- max_x   |     | image
      min_y                               min_y   |     | height
      |                                       |   |     |
      |                                       |   |     v
      |                                       |   |
      |                                       |   / max_x, max_y,     ^
      |                                       |  /  max_depth        /
      min_x                               max_y /                   / n_pts_per_ray
      max_y ----------------------------- max_x/ min_depth         v
              < --- image_width --- >
  
  In order to generate ray points, `MultinomialRaysampler` takes each 3D point of
  the grid (with coordinates `[x, y, depth]`) and unprojects it
  with `cameras.unproject_points([x, y, depth])`, where `cameras` are an
  additional input to the `forward` function.
  
  Note that this is a generic implementation that can support any image grid
  coordinate convention. For a raysampler which follows the PyTorch3D
  coordinate conventions please refer to `NDCMultinomialRaysampler`.
  As such, `NDCMultinomialRaysampler` is a special case of `MultinomialRaysampler`.
  
  Attributes:
      min_x: The leftmost x-coordinate of each ray's source pixel's center.
      max_x: The rightmost x-coordinate of each ray's source pixel's center.
      min_y: The topmost y-coordinate of each ray's source pixel's center.
      max_y: The bottommost y-coordinate of each ray's source pixel's center.
class_docstr_tok: ["Samples","a","fixed","number","of","points","along","rays","which","are","regularly","distributed","in","a","batch","of","rectangular","image","grids",".","Points","along","each","ray","have","uniformly-spaced","z-coordinates","between","a","predefined","minimum","and","maximum","depth",".","The","raysampler","first","generates","a","3D","coordinate","grid","of","the","following","form",":",":","\/","min_x",",","min_y",",","max_depth","--","--","--","--","--","--","--","\/","max_x",",","min_y",",","max_depth","\/","\/|","\/","\/","|","^","\/","min_depth","min_depth","\/","|","|","min_x","--","--","--","--","--","--","--","--","--","--","--","--","--","--","-","max_x","|","|","image","min_y","min_y","|","|","height","|","|","|","|","|","|","|","v","|","|","|","|","|","\/","max_x",",","max_y",",","^","|","|","\/","max_depth","\/","min_x","max_y","\/","\/","n_pts_per_ray","max_y","--","--","--","--","--","--","--","--","--","--","--","--","--","--","-","max_x\/","min_depth","v","<","--","-","image_width","--","-",">","In","order","to","generate","ray","points",",","`","MultinomialRaysampler","`","takes","each","3D","point","of","the","grid","(","with","coordinates","`","[","x",",","y",",","depth","]","`",")","and","unprojects","it","with","`","cameras.unproject_points","(","[","x",",","y",",","depth","]",")","`",",","where","`","cameras","`","are","an","additional","input","to","the","`","forward","`","function",".","Note","that","this","is","a","generic","implementation","that","can","support","any","image","grid","coordinate","convention",".","For","a","raysampler","which","follows","the","PyTorch3D","coordinate","conventions","please","refer","to","`","NDCMultinomialRaysampler","`",".","As","such",",","`","NDCMultinomialRaysampler","`","is","a","special","case","of","`","MultinomialRaysampler","`",".","Attributes",":","min_x",":","The","leftmost","x-coordinate","of","each","ray","'s","source","pixel","'s","center",".","max_x",":","The","rightmost","x-coordinate","of","each","ray","'s","source","pixel","'s","center",".","min_y",":","The","topmost","y-coordinate","of","each","ray","'s","source","pixel","'s","center",".","max_y",":","The","bottommost","y-coordinate","of","each","ray","'s","source","pixel","'s","center","."]
docstr: Args:
      cameras: A batch of `batch_size` cameras from which the rays are emitted.
      mask: if given, the rays are sampled from the mask. Should be of size
          (batch_size, image_height, image_width).
      min_depth: The minimum depth of a ray-point.
      max_depth: The maximum depth of a ray-point.
      n_rays_per_image: If given, this amount of rays are sampled from the grid.
          `n_rays_per_image` and `n_rays_total` cannot both be defined.
      n_pts_per_ray: The number of points sampled along each ray.
      stratified_sampling: if set, overrides stratified_sampling provided
          in __init__.
      n_rays_total: How many rays in total to sample from the cameras provided. The result
          is as if `n_rays_total_training` cameras were sampled with replacement from the
          cameras provided and for every camera one ray was sampled. If set returns the
          HeterogeneousRayBundle with batch_size=n_rays_total.
          `n_rays_per_image` and `n_rays_total` cannot both be defined.
  Returns:
      A named tuple RayBundle or dataclass HeterogeneousRayBundle with the
      following fields:
  
      origins: A tensor of shape
          `(batch_size, s1, s2, 3)`
          denoting the locations of ray origins in the world coordinates.
      directions: A tensor of shape
          `(batch_size, s1, s2, 3)`
          denoting the directions of each ray in the world coordinates.
      lengths: A tensor of shape
          `(batch_size, s1, s2, n_pts_per_ray)`
          containing the z-coordinate (=depth) of each ray in world units.
      xys: A tensor of shape
          `(batch_size, s1, s2, 2)`
          containing the 2D image coordinates of each ray or,
          if mask is given, `(batch_size, n, 1, 2)`
      Here `s1, s2` refer to spatial dimensions.
      `(s1, s2)` refer to (highest priority first):
          - `(1, 1)` if `n_rays_total` is provided, (batch_size=n_rays_total)
          - `(n_rays_per_image, 1) if `n_rays_per_image` if provided,
          - `(n, 1)` where n is the minimum cardinality of the mask
                  in the batch if `mask` is provided
          - `(image_height, image_width)` if nothing from above is satisfied
  
      `HeterogeneousRayBundle` has additional members:
          - camera_ids: tensor of shape (M,), where `M` is the number of unique sampled
              cameras. It represents unique ids of sampled cameras.
          - camera_counts: tensor of shape (M,), where `M` is the number of unique sampled
              cameras. Represents how many times each camera from `camera_ids` was sampled
  
      `HeterogeneousRayBundle` is returned if `n_rays_total` is provided else `RayBundle`
      is returned.
docstr_tok: ["Args",":","cameras",":","A","batch","of","`","batch_size","`","cameras","from","which","the","rays","are","emitted",".","mask",":","if","given",",","the","rays","are","sampled","from","the","mask",".","Should","be","of","size","(","batch_size",",","image_height",",","image_width",")",".","min_depth",":","The","minimum","depth","of","a","ray-point",".","max_depth",":","The","maximum","depth","of","a","ray-point",".","n_rays_per_image",":","If","given",",","this","amount","of","rays","are","sampled","from","the","grid",".","`","n_rays_per_image","`","and","`","n_rays_total","`","can","not","both","be","defined",".","n_pts_per_ray",":","The","number","of","points","sampled","along","each","ray",".","stratified_sampling",":","if","set",",","overrides","stratified_sampling","provided","in","__init__",".","n_rays_total",":","How","many","rays","in","total","to","sample","from","the","cameras","provided",".","The","result","is","as","if","`","n_rays_total_training","`","cameras","were","sampled","with","replacement","from","the","cameras","provided","and","for","every","camera","one","ray","was","sampled",".","If","set","returns","the","HeterogeneousRayBundle","with","batch_size=n_rays_total",".","`","n_rays_per_image","`","and","`","n_rays_total","`","can","not","both","be","defined",".","Returns",":","A","named","tuple","RayBundle","or","dataclass","HeterogeneousRayBundle","with","the","following","fields",":","origins",":","A","tensor","of","shape","`","(","batch_size",",","s1",",","s2",",","3",")","`","denoting","the","locations","of","ray","origins","in","the","world","coordinates",".","directions",":","A","tensor","of","shape","`","(","batch_size",",","s1",",","s2",",","3",")","`","denoting","the","directions","of","each","ray","in","the","world","coordinates",".","lengths",":","A","tensor","of","shape","`","(","batch_size",",","s1",",","s2",",","n_pts_per_ray",")","`","containing","the","z-coordinate","(","=depth",")","of","each","ray","in","world","units",".","xys",":","A","tensor","of","shape","`","(","batch_size",",","s1",",","s2",",","2",")","`","containing","the","2D","image","coordinates","of","each","ray","or",",","if","mask","is","given",",","`","(","batch_size",",","n",",","1",",","2",")","`","Here","`","s1",",","s2","`","refer","to","spatial","dimensions",".","`","(","s1",",","s2",")","`","refer","to","(","highest","priority","first",")",":","-","`","(","1",",","1",")","`","if","`","n_rays_total","`","is","provided",",","(","batch_size=n_rays_total",")","-","`","(","n_rays_per_image",",","1",")","if","`","n_rays_per_image","`","if","provided",",","-","`","(","n",",","1",")","`","where","n","is","the","minimum","cardinality","of","the","mask","in","the","batch","if","`","mask","`","is","provided","-","`","(","image_height",",","image_width",")","`","if","nothing","from","above","is","satisfied","`","HeterogeneousRayBundle","`","has","additional","members",":","-","camera_ids",":","tensor","of","shape","(","M",",",")",",","where","`","M","`","is","the","number","of","unique","sampled","cameras",".","It","represents","unique","ids","of","sampled","cameras",".","-","camera_counts",":","tensor","of","shape","(","M",",",")",",","where","`","M","`","is","the","number","of","unique","sampled","cameras",".","Represents","how","many","times","each","camera","from","`","camera_ids","`","was","sampled","`","HeterogeneousRayBundle","`","is","returned","if","`","n_rays_total","`","is","provided","else","`","RayBundle","`","is","returned","."]
code: def forward(
      self,
      cameras: CamerasBase,
      *,
      mask: Optional[torch.Tensor] = None,
      min_depth: Optional[float] = None,
      max_depth: Optional[float] = None,
      n_rays_per_image: Optional[int] = None,
      n_pts_per_ray: Optional[int] = None,
      stratified_sampling: Optional[bool] = None,
      n_rays_total: Optional[int] = None,
      **kwargs,
  ) -> Union[RayBundle, HeterogeneousRayBundle]:
      """
      Args:
          cameras: A batch of `batch_size` cameras from which the rays are emitted.
          mask: if given, the rays are sampled from the mask. Should be of size
              (batch_size, image_height, image_width).
          min_depth: The minimum depth of a ray-point.
          max_depth: The maximum depth of a ray-point.
          n_rays_per_image: If given, this amount of rays are sampled from the grid.
              `n_rays_per_image` and `n_rays_total` cannot both be defined.
          n_pts_per_ray: The number of points sampled along each ray.
          stratified_sampling: if set, overrides stratified_sampling provided
              in __init__.
          n_rays_total: How many rays in total to sample from the cameras provided. The result
              is as if `n_rays_total_training` cameras were sampled with replacement from the
              cameras provided and for every camera one ray was sampled. If set returns the
              HeterogeneousRayBundle with batch_size=n_rays_total.
              `n_rays_per_image` and `n_rays_total` cannot both be defined.
      Returns:
          A named tuple RayBundle or dataclass HeterogeneousRayBundle with the
          following fields:
  
          origins: A tensor of shape
              `(batch_size, s1, s2, 3)`
              denoting the locations of ray origins in the world coordinates.
          directions: A tensor of shape
              `(batch_size, s1, s2, 3)`
              denoting the directions of each ray in the world coordinates.
          lengths: A tensor of shape
              `(batch_size, s1, s2, n_pts_per_ray)`
              containing the z-coordinate (=depth) of each ray in world units.
          xys: A tensor of shape
              `(batch_size, s1, s2, 2)`
              containing the 2D image coordinates of each ray or,
              if mask is given, `(batch_size, n, 1, 2)`
          Here `s1, s2` refer to spatial dimensions.
          `(s1, s2)` refer to (highest priority first):
              - `(1, 1)` if `n_rays_total` is provided, (batch_size=n_rays_total)
              - `(n_rays_per_image, 1) if `n_rays_per_image` if provided,
              - `(n, 1)` where n is the minimum cardinality of the mask
                      in the batch if `mask` is provided
              - `(image_height, image_width)` if nothing from above is satisfied
  
          `HeterogeneousRayBundle` has additional members:
              - camera_ids: tensor of shape (M,), where `M` is the number of unique sampled
                  cameras. It represents unique ids of sampled cameras.
              - camera_counts: tensor of shape (M,), where `M` is the number of unique sampled
                  cameras. Represents how many times each camera from `camera_ids` was sampled
  
          `HeterogeneousRayBundle` is returned if `n_rays_total` is provided else `RayBundle`
          is returned.
      """
      n_rays_total = n_rays_total or self._n_rays_total
      n_rays_per_image = n_rays_per_image or self._n_rays_per_image
      if (n_rays_total is not None) and (n_rays_per_image is not None):
          raise ValueError(
              "`n_rays_total` and `n_rays_per_image` cannot both be defined."
          )
      if n_rays_total:
          (
              cameras,
              mask,
              camera_ids,  # unique ids of sampled cameras
              camera_counts,  # number of times unique camera id was sampled
              # `n_rays_per_image` is equal to the max number of times a simgle camera
              # was sampled. We sample all cameras at `camera_ids` `n_rays_per_image` times
              # and then discard the unneeded rays.
              # pyre-ignore[9]
              n_rays_per_image,
          ) = _sample_cameras_and_masks(n_rays_total, cameras, mask)
      else:
          # pyre-ignore[9]
          camera_ids: torch.LongTensor = torch.arange(
              len(cameras), dtype=torch.long
          )
  
      batch_size = cameras.R.shape[0]
      device = cameras.device
  
      # expand the (H, W, 2) grid batch_size-times to (B, H, W, 2)
      xy_grid = self._xy_grid.to(device).expand(batch_size, -1, -1, -1)
  
      if mask is not None and n_rays_per_image is None:
          # if num rays not given, sample according to the smallest mask
          n_rays_per_image = (
              n_rays_per_image
              or mask.sum(dim=(1, 2)).min().int().item()
          )
  
      if n_rays_per_image is not None:
          if mask is not None:
              assert mask.shape == xy_grid.shape[:3]
              weights = mask.reshape(batch_size, -1)
          else:
              # it is probably more efficient to use torch.randperm
              # for uniform weights but it is unlikely given that randperm
              # is not batched and does not support partial permutation
              _, width, height, _ = xy_grid.shape
              weights = xy_grid.new_ones(batch_size, width * height)
          # pyre-fixme[6]: For 2nd param expected `int` but got `Union[bool,
          #  float, int]`.
          rays_idx = _safe_multinomial(weights, n_rays_per_image)[
              ..., None
          ].expand(-1, -1, 2)
  
          xy_grid = torch.gather(
              xy_grid.reshape(batch_size, -1, 2), 1, rays_idx
          )[:, :, None]
  
      min_depth = (
          min_depth if min_depth is not None else self._min_depth
      )
      max_depth = (
          max_depth if max_depth is not None else self._max_depth
      )
      n_pts_per_ray = (
          n_pts_per_ray
          if n_pts_per_ray is not None
          else self._n_pts_per_ray
      )
      stratified_sampling = (
          stratified_sampling
          if stratified_sampling is not None
          else self._stratified_sampling
      )
  
      ray_bundle = _xy_to_ray_bundle(
          cameras,
          xy_grid,
          min_depth,
          max_depth,
          n_pts_per_ray,
          self._unit_directions,
          stratified_sampling,
      )
  
      return (
          # pyre-ignore[61]
          _pack_ray_bundle(ray_bundle, camera_ids, camera_counts)
          if n_rays_total
          else ray_bundle
      )
  
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 2 for calling the MultinomialRaysampler obj's underlying member method code with example usage: obj.forward(cameras) and returns: unknown