class_docstr: Samples a fixed number of points along rays which are in turn sampled for
  each camera in a batch.
  
  This class utilizes `NDCMultinomialRaysampler` which allows to either
  randomly sample rays from an input foreground saliency mask
  (`RenderSamplingMode.MASK_SAMPLE`), or on a rectangular image grid
  (`RenderSamplingMode.FULL_GRID`). The sampling mode can be set separately
  for training and evaluation by setting `self.sampling_mode_training`
  and `self.sampling_mode_training` accordingly.
  
  The class allows to adjust the sampling points along rays by overwriting the
  `AbstractMaskRaySampler._get_min_max_depth_bounds` function which returns
  the near/far planes (`min_depth`/`max_depth`) `NDCMultinomialRaysampler`.
  
  Settings:
      image_width: The horizontal size of the image grid.
      image_height: The vertical size of the image grid.
      sampling_mode_training: The ray sampling mode for training. This should be a str
          option from the RenderSamplingMode Enum
      sampling_mode_evaluation: Same as above but for evaluation.
      n_pts_per_ray_training: The number of points sampled along each ray during training.
      n_pts_per_ray_evaluation: The number of points sampled along each ray during evaluation.
      n_rays_per_image_sampled_from_mask: The amount of rays to be sampled from the image
          grid. Given a batch of image grids, this many is sampled from each.
          `n_rays_per_image_sampled_from_mask` and `n_rays_total_training` cannot both be
          defined.
      n_rays_total_training: (optional) How many rays in total to sample from the entire
          batch of provided image grid. The result is as if `n_rays_total_training`
          cameras/image grids were sampled with replacement from the cameras / image grids
          provided and for every camera one ray was sampled.
          `n_rays_per_image_sampled_from_mask` and `n_rays_total_training` cannot both be
          defined, to use you have to set `n_rays_per_image` to None.
          Used only for EvaluationMode.TRAINING.
      stratified_point_sampling_training: if set, performs stratified random sampling
          along the ray; otherwise takes ray points at deterministic offsets.
      stratified_point_sampling_evaluation: Same as above but for evaluation.
      cast_ray_bundle_as_cone: If True, the sampling will generate the bins and radii
          attribute of ImplicitronRayBundle. The `bins` contain the z-coordinate
          (=depth) of each ray in world units and are of shape
          `(batch_size, n_rays_per_image, n_pts_per_ray_training/evaluation + 1)`
          while `lengths` is equal to the midpoint of the bins:
          (0.5 * (bins[..., 1:] + bins[..., :-1]).
          If False, `bins` is None, `radii` is None and `lengths` contains
          the z-coordinate (=depth) of each ray in world units and are of shape
          `(batch_size, n_rays_per_image, n_pts_per_ray_training/evaluation)`
  
  Raises:
      TypeError: if cast_ray_bundle_as_cone is set to True and n_rays_total_training
          is not None will result in an error. HeterogeneousRayBundle is
          not supported for conical frustum computation yet.
docstr: Args:
      cameras: A batch of `batch_size` cameras from which the rays are emitted.
      evaluation_mode: one of `EvaluationMode.TRAINING` or
          `EvaluationMode.EVALUATION` which determines the sampling mode
          that is used.
      mask: Active for the `RenderSamplingMode.MASK_SAMPLE` sampling mode.
          Defines a non-negative mask of shape
          `(batch_size, image_height, image_width)` where each per-pixel
          value is proportional to the probability of sampling the
          corresponding pixel's ray.
  
  Returns:
      ray_bundle: A `ImplicitronRayBundle` object containing the parametrizations of the
          sampled rendering rays.
docstr_tok: ["Args",":","cameras",":","A","batch","of","`","batch_size","`","cameras","from","which","the","rays","are","emitted",".","evaluation_mode",":","one","of","`","EvaluationMode.TRAINING","`","or","`","EvaluationMode.EVALUATION","`","which","determines","the","sampling","mode","that","is","used",".","mask",":","Active","for","the","`","RenderSamplingMode.MASK_SAMPLE","`","sampling","mode",".","Defines","a","non-negative","mask","of","shape","`","(","batch_size",",","image_height",",","image_width",")","`","where","each","per-pixel","value","is","proportional","to","the","probability","of","sampling","the","corresponding","pixel","'s","ray",".","Returns",":","ray_bundle",":","A","`","ImplicitronRayBundle","`","object","containing","the","parametrizations","of","the","sampled","rendering","rays","."]
code: def forward(
      self,
      cameras: CamerasBase,
      evaluation_mode: EvaluationMode,
      mask: Optional[torch.Tensor] = None,
  ) -> ImplicitronRayBundle:
      """
  
      Args:
          cameras: A batch of `batch_size` cameras from which the rays are emitted.
          evaluation_mode: one of `EvaluationMode.TRAINING` or
              `EvaluationMode.EVALUATION` which determines the sampling mode
              that is used.
          mask: Active for the `RenderSamplingMode.MASK_SAMPLE` sampling mode.
              Defines a non-negative mask of shape
              `(batch_size, image_height, image_width)` where each per-pixel
              value is proportional to the probability of sampling the
              corresponding pixel's ray.
  
      Returns:
          ray_bundle: A `ImplicitronRayBundle` object containing the parametrizations of the
              sampled rendering rays.
      """
      sample_mask = None
      if (
          self._sampling_mode[evaluation_mode]
          == RenderSamplingMode.MASK_SAMPLE
          and mask is not None
      ):
          sample_mask = torch.nn.functional.interpolate(
              mask,
              size=[self.image_height, self.image_width],
              mode="nearest",
          )[:, 0]
  
      min_depth, max_depth = self._get_min_max_depth_bounds(cameras)
  
      raysampler = {
          EvaluationMode.TRAINING: self._training_raysampler,
          EvaluationMode.EVALUATION: self._evaluation_raysampler,
      }[evaluation_mode]
  
      ray_bundle = raysampler(
          cameras=cameras,
          mask=sample_mask,
          min_depth=min_depth,
          max_depth=max_depth,
      )
      if self.cast_ray_bundle_as_cone and isinstance(
          ray_bundle, HeterogeneousRayBundle
      ):
          # If this error rises it means that raysampler has among
          # its arguments `n_ray_totals`. If it is the case
          # then you should update the radii computation and lengths
          # computation to handle padding and unpadding.
          raise TypeError(
              "Heterogeneous ray bundle is not supported for conical frustum computation yet"
          )
      elif self.cast_ray_bundle_as_cone:
          pixel_hw: Tuple[float, float] = (
              self.pixel_height,
              self.pixel_width,
          )
          pixel_radii_2d = compute_radii(
              cameras, ray_bundle.xys[..., :2], pixel_hw
          )
          return ImplicitronRayBundle(
              directions=ray_bundle.directions,
              origins=ray_bundle.origins,
              lengths=None,
              xys=ray_bundle.xys,
              bins=ray_bundle.lengths,
              pixel_radii_2d=pixel_radii_2d,
          )
  
      return ImplicitronRayBundle(
          directions=ray_bundle.directions,
          origins=ray_bundle.origins,
          lengths=ray_bundle.lengths,
          xys=ray_bundle.xys,
          camera_counts=getattr(ray_bundle, "camera_counts", None),
          camera_ids=getattr(ray_bundle, "camera_ids", None),
      )
  
imports: from typing import Optional, Tuple
  import torch
  from pytorch3d.implicitron.tools import camera_utils
  from pytorch3d.implicitron.tools.config import registry, ReplaceableBase
  from pytorch3d.renderer import NDCMultinomialRaysampler
  from pytorch3d.renderer.cameras import CamerasBase
  from pytorch3d.renderer.implicit.utils import HeterogeneousRayBundle
  from .base import EvaluationMode, ImplicitronRayBundle, RenderSamplingMode
label_desc: Use image node_id 3 for calling the AbstractMaskRaySampler obj's underlying member method code with example usage: obj.forward(cameras, evaluation_mode, mask) and returns: ImplicitronRayBundle, ImplicitronRayBundle