class_docstr: Implements the raysampler of NeRF.
  
  Depending on the `self.training` flag, the raysampler either samples
  a chunk of random rays (`self.training==True`), or returns a subset of rays
  of the full image grid (`self.training==False`).
  The chunking of rays allows for efficient evaluation of the NeRF implicit
  surface function without encountering out-of-GPU-memory errors.
  
  Additionally, this raysampler supports pre-caching of the ray bundles
  for a set of input cameras (`self.precache_rays`).
  Pre-caching the rays before training greatly speeds-up the ensuing
  raysampling step of the training NeRF iterations.
class_docstr_tok: ["Implements","the","raysampler","of","NeRF",".","Depending","on","the","`","self.training","`","flag",",","the","raysampler","either","samples","a","chunk","of","random","rays","(","`","self.training==True","`",")",",","or","returns","a","subset","of","rays","of","the","full","image","grid","(","`","self.training==False","`",")",".","The","chunking","of","rays","allows","for","efficient","evaluation","of","the","NeRF","implicit","surface","function","without","encountering","out-of-GPU-memory","errors",".","Additionally",",","this","raysampler","supports","pre-caching","of","the","ray","bundles","for","a","set","of","input","cameras","(","`","self.precache_rays","`",")",".","Pre-caching","the","rays","before","training","greatly","speeds-up","the","ensuing","raysampling","step","of","the","training","NeRF","iterations","."]
docstr: Args:
      cameras: A batch of `batch_size` cameras from which the rays are emitted.
      chunksize: The number of rays per chunk.
          Active only when `self.training==False`.
      chunk_idx: The index of the ray chunk. The number has to be in
          `[0, self.get_n_chunks(chunksize, batch_size)-1]`.
          Active only when `self.training==False`.
      camera_hash: A unique identifier of a pre-cached camera. If `None`,
          the cache is not searched and the rays are calculated from scratch.
      caching: If `True`, activates the caching mode that returns the `RayBundle`
          that should be stored into the cache.
  Returns:
      A named tuple `RayBundle` with the following fields:
          origins: A tensor of shape
              `(batch_size, n_rays_per_image, 3)`
              denoting the locations of ray origins in the world coordinates.
          directions: A tensor of shape
              `(batch_size, n_rays_per_image, 3)`
              denoting the directions of each ray in the world coordinates.
          lengths: A tensor of shape
              `(batch_size, n_rays_per_image, n_pts_per_ray)`
              containing the z-coordinate (=depth) of each ray in world units.
          xys: A tensor of shape
              `(batch_size, n_rays_per_image, 2)`
              containing the 2D image coordinates of each ray.
docstr_tok: ["Args",":","cameras",":","A","batch","of","`","batch_size","`","cameras","from","which","the","rays","are","emitted",".","chunksize",":","The","number","of","rays","per","chunk",".","Active","only","when","`","self.training==False","`",".","chunk_idx",":","The","index","of","the","ray","chunk",".","The","number","has","to","be","in","`","[","0",",","self.get_n_chunks","(","chunksize",",","batch_size",")","-1","]","`",".","Active","only","when","`","self.training==False","`",".","camera_hash",":","A","unique","identifier","of","a","pre-cached","camera",".","If","`","None","`",",","the","cache","is","not","searched","and","the","rays","are","calculated","from","scratch",".","caching",":","If","`","True","`",",","activates","the","caching","mode","that","returns","the","`","RayBundle","`","that","should","be","stored","into","the","cache",".","Returns",":","A","named","tuple","`","RayBundle","`","with","the","following","fields",":","origins",":","A","tensor","of","shape","`","(","batch_size",",","n_rays_per_image",",","3",")","`","denoting","the","locations","of","ray","origins","in","the","world","coordinates",".","directions",":","A","tensor","of","shape","`","(","batch_size",",","n_rays_per_image",",","3",")","`","denoting","the","directions","of","each","ray","in","the","world","coordinates",".","lengths",":","A","tensor","of","shape","`","(","batch_size",",","n_rays_per_image",",","n_pts_per_ray",")","`","containing","the","z-coordinate","(","=depth",")","of","each","ray","in","world","units",".","xys",":","A","tensor","of","shape","`","(","batch_size",",","n_rays_per_image",",","2",")","`","containing","the","2D","image","coordinates","of","each","ray","."]
code: def forward(
      self,
      cameras: CamerasBase,
      chunksize: int = None,
      chunk_idx: int = 0,
      camera_hash: str = None,
      caching: bool = False,
      **kwargs,
  ) -> RayBundle:
      """
      Args:
          cameras: A batch of `batch_size` cameras from which the rays are emitted.
          chunksize: The number of rays per chunk.
              Active only when `self.training==False`.
          chunk_idx: The index of the ray chunk. The number has to be in
              `[0, self.get_n_chunks(chunksize, batch_size)-1]`.
              Active only when `self.training==False`.
          camera_hash: A unique identifier of a pre-cached camera. If `None`,
              the cache is not searched and the rays are calculated from scratch.
          caching: If `True`, activates the caching mode that returns the `RayBundle`
              that should be stored into the cache.
      Returns:
          A named tuple `RayBundle` with the following fields:
              origins: A tensor of shape
                  `(batch_size, n_rays_per_image, 3)`
                  denoting the locations of ray origins in the world coordinates.
              directions: A tensor of shape
                  `(batch_size, n_rays_per_image, 3)`
                  denoting the directions of each ray in the world coordinates.
              lengths: A tensor of shape
                  `(batch_size, n_rays_per_image, n_pts_per_ray)`
                  containing the z-coordinate (=depth) of each ray in world units.
              xys: A tensor of shape
                  `(batch_size, n_rays_per_image, 2)`
                  containing the 2D image coordinates of each ray.
      """
  
      batch_size = cameras.R.shape[0]  # pyre-ignore
      device = cameras.device
  
      if (camera_hash is None) and (not caching) and self.training:
          # Sample random rays from scratch.
          ray_bundle = self._mc_raysampler(cameras)
          ray_bundle = self._normalize_raybundle(ray_bundle)
      else:
          if camera_hash is not None:
              # The case where we retrieve a camera from cache.
              if batch_size != 1:
                  raise NotImplementedError(
                      "Ray caching works only for batches with a single camera!"
                  )
              full_ray_bundle = self._ray_cache[camera_hash]
          else:
              # We generate a full ray grid from scratch.
              full_ray_bundle = self._grid_raysampler(cameras)
              full_ray_bundle = self._normalize_raybundle(
                  full_ray_bundle
              )
  
          n_pixels = full_ray_bundle.directions.shape[:-1].numel()
  
          if self.training:
              # During training we randomly subsample rays.
              sel_rays = torch.randperm(
                  n_pixels, device=full_ray_bundle.lengths.device
              )[: self._mc_raysampler._n_rays_per_image]
          else:
              # In case we test, we take only the requested chunk.
              if chunksize is None:
                  chunksize = n_pixels * batch_size
              start = chunk_idx * chunksize * batch_size
              end = min(start + chunksize, n_pixels)
              sel_rays = torch.arange(
                  start,
                  end,
                  dtype=torch.long,
                  device=full_ray_bundle.lengths.device,
              )
  
          # Take the "sel_rays" rays from the full ray bundle.
          ray_bundle = RayBundle(
              *[
                  v.view(n_pixels, -1)[sel_rays]
                  .view(batch_size, sel_rays.numel() // batch_size, -1)
                  .to(device)
                  for v in full_ray_bundle
              ]
          )
  
      if (
          (self._stratified and self.training)
          or (self._stratified_test and not self.training)
      ) and not caching:  # Make sure not to stratify when caching!
          ray_bundle = self._stratify_ray_bundle(ray_bundle)
  
      return ray_bundle
  
file_path: pytorch3d/projects/nerf/nerf/raysampler.py
imports: import math
  from typing import List
  import torch
  from pytorch3d.renderer import MonteCarloRaysampler, NDCMultinomialRaysampler, RayBundle
  from pytorch3d.renderer.cameras import CamerasBase
  from pytorch3d.renderer.implicit.sample_pdf import sample_pdf
label_desc: Use image node_id 7 for calling the NeRFRaysampler obj's underlying member method code with example usage: obj.forward(cameras, chunksize, chunk_idx, camera_hash, caching) and returns: ray_bundle