class_docstr: This implicit function consists of two streams, one for the density calculation and one
  for the color calculation. Each of these streams has three main parts:
      1) Voxel grids:
          They take the (x, y, z) position and return the embedding of that point.
          These components are replaceable, you can make your own or choose one of
          several options.
      2) Harmonic embeddings:
          Convert each feature into series of 'harmonic features', feature is passed through
          sine and cosine functions. Input is of shape [minibatch, ..., D] output
          [minibatch, ..., (n_harmonic_functions * 2 + int(append_input)) * D]. Appends
          input by default. If you want it to behave like identity, put n_harmonic_functions=0
          and append_input=True.
      3) Decoding functions:
          The decoder is an instance of the DecoderFunctionBase and converts the embedding
          of a spatial location to density/color. Examples are Identity which returns its
          input and the MLP which uses fully connected nerual network to transform the input.
          These components are replaceable, you can make your own or choose from
          several options.
  
  Calculating density is done in three steps:
      1) Evaluating the voxel grid on points
      2) Embedding the outputs with harmonic embedding
      3) Passing through the Density decoder
  
  To calculate the color we need the embedding and the viewing direction, it has five steps:
      1) Transforming the viewing direction with camera
      2) Evaluating the voxel grid on points
      3) Embedding the outputs with harmonic embedding
      4) Embedding the normalized direction with harmonic embedding
      5) Passing everything through the Color decoder
  
  If using the Implicitron configuration system the input_dim to the decoding functions will
  be set to the output_dim of the Harmonic embeddings.
  
  A speed up comes from using the scaffold, a low resolution voxel grid.
  The scaffold is referenced as "binary occupancy grid mask" in TensoRF paper and "AlphaMask"
  in official TensoRF implementation.
  The scaffold is used in:
      1) filtering points in empty space
          - controlled by `scaffold_filter_points` boolean. If set to True, points for which
              scaffold predicts that are in empty space will return 0 density and
              (0, 0, 0) color.
      2) calculating the bounding box of an object and cropping the voxel grids
          - controlled by `volume_cropping_epochs`.
          - at those epochs the implicit function will find the bounding box of an object
              inside it and crop density and color grids. Cropping of the voxel grids means
              preserving only voxel values that are inside the bounding box and changing the
              resolution to match the original, while preserving the new cropped location in
              world coordinates.
  
  The scaffold has to exist before attempting filtering and cropping, and is created on
  `scaffold_calculating_epochs`. Each voxel in the scaffold is labeled as having density 1 if
  the point in the center of it evaluates to greater than `scaffold_empty_space_threshold`.
  3D max pooling is performed on the densities of the points in 3D.
  Scaffold features are off by default.
  
  Members:
      voxel_grid_density (VoxelGridBase): voxel grid to use for density estimation
      voxel_grid_color   (VoxelGridBase): voxel grid to use for color   estimation
  
      harmonic_embedder_xyz_density (HarmonicEmbedder): Function to transform the outputs of
          the voxel_grid_density
      harmonic_embedder_xyz_color (HarmonicEmbedder): Function to transform the outputs of
          the voxel_grid_color for density
      harmonic_embedder_dir_color (HarmonicEmbedder): Function to transform the outputs of
          the voxel_grid_color for color
  
      decoder_density (DecoderFunctionBase): decoder function to use for density estimation
      color_density   (DecoderFunctionBase): decoder function to use for color   estimation
  
      use_multiple_streams (bool): if you want the density and color calculations to run on
          different cuda streams set this to True. Default True.
      xyz_ray_dir_in_camera_coords (bool): This is true if the directions are given in
          camera coordinates. Default False.
  
      voxel_grid_scaffold (VoxelGridModule): which holds the scaffold. Extents and
          translation of it are set to those of voxel_grid_density.
      scaffold_calculating_epochs (Tuple[int, ...]): at which epochs to recalculate the
          scaffold. (The scaffold will be created automatically at the beginning of
          the calculation.)
      scaffold_resolution (Tuple[int, int, int]): (width, height, depth) of the underlying
          voxel grid which stores scaffold
      scaffold_empty_space_threshold (float): if `self._get_density` evaluates to less than
          this it will be considered as empty space and the scaffold at that point would
          evaluate as empty space.
      scaffold_occupancy_chunk_size (str or int): Number of xy scaffold planes to calculate
          at the same time. To calculate the scaffold we need to query `_get_density()` at
          every voxel, this calculation can be split into scaffold depth number of xy plane
          calculations if you want the lowest memory usage, one calculation to calculate the
          whole scaffold, but with higher memory footprint or any other number of planes.
          Setting to a non-positive number calculates all planes at the same time.
          Defaults to '-1' (=calculating all planes).
      scaffold_max_pool_kernel_size (int): Size of the pooling region to use when
          calculating the scaffold. Defaults to 3.
      scaffold_filter_points (bool): If set to True the points will be filtered using
          `self.voxel_grid_scaffold`. Filtered points will be predicted as having 0 density
          and (0, 0, 0) color. The points which were not evaluated as empty space will be
          passed through the steps outlined above.
      volume_cropping_epochs: on which epochs to crop the voxel grids to fit the object's
          bounding box. Scaffold has to be calculated before cropping.
class_docstr_tok: ["This","implicit","function","consists","of","two","streams",",","one","for","the","density","calculation","and","one","for","the","color","calculation",".","Each","of","these","streams","has","three","main","parts",":","1",")","Voxel","grids",":","They","take","the","(","x",",","y",",","z",")","position","and","return","the","embedding","of","that","point",".","These","components","are","replaceable",",","you","can","make","your","own","or","choose","one","of","several","options",".","2",")","Harmonic","embeddings",":","Convert","each","feature","into","series","of","'harmonic","features","'",",","feature","is","passed","through","sine","and","cosine","functions",".","Input","is","of","shape","[","minibatch",",","...",",","D","]","output","[","minibatch",",","...",",","(","n_harmonic_functions","*","2","+","int","(","append_input",")",")","*","D","]",".","Appends","input","by","default",".","If","you","want","it","to","behave","like","identity",",","put","n_harmonic_functions=0","and","append_input=True",".","3",")","Decoding","functions",":","The","decoder","is","an","instance","of","the","DecoderFunctionBase","and","converts","the","embedding","of","a","spatial","location","to","density\/color",".","Examples","are","Identity","which","returns","its","input","and","the","MLP","which","uses","fully","connected","nerual","network","to","transform","the","input",".","These","components","are","replaceable",",","you","can","make","your","own","or","choose","from","several","options",".","Calculating","density","is","done","in","three","steps",":","1",")","Evaluating","the","voxel","grid","on","points","2",")","Embedding","the","outputs","with","harmonic","embedding","3",")","Passing","through","the","Density","decoder","To","calculate","the","color","we","need","the","embedding","and","the","viewing","direction",",","it","has","five","steps",":","1",")","Transforming","the","viewing","direction","with","camera","2",")","Evaluating","the","voxel","grid","on","points","3",")","Embedding","the","outputs","with","harmonic","embedding","4",")","Embedding","the","normalized","direction","with","harmonic","embedding","5",")","Passing","everything","through","the","Color","decoder","If","using","the","Implicitron","configuration","system","the","input_dim","to","the","decoding","functions","will","be","set","to","the","output_dim","of","the","Harmonic","embeddings",".","A","speed","up","comes","from","using","the","scaffold",",","a","low","resolution","voxel","grid",".","The","scaffold","is","referenced","as","``","binary","occupancy","grid","mask","''","in","TensoRF","paper","and","``","AlphaMask","''","in","official","TensoRF","implementation",".","The","scaffold","is","used","in",":","1",")","filtering","points","in","empty","space","-","controlled","by","`","scaffold_filter_points","`","boolean",".","If","set","to","True",",","points","for","which","scaffold","predicts","that","are","in","empty","space","will","return","0","density","and","(","0",",","0",",","0",")","color",".","2",")","calculating","the","bounding","box","of","an","object","and","cropping","the","voxel","grids","-","controlled","by","`","volume_cropping_epochs","`",".","-","at","those","epochs","the","implicit","function","will","find","the","bounding","box","of","an","object","inside","it","and","crop","density","and","color","grids",".","Cropping","of","the","voxel","grids","means","preserving","only","voxel","values","that","are","inside","the","bounding","box","and","changing","the","resolution","to","match","the","original",",","while","preserving","the","new","cropped","location","in","world","coordinates",".","The","scaffold","has","to","exist","before","attempting","filtering","and","cropping",",","and","is","created","on","`","scaffold_calculating_epochs","`",".","Each","voxel","in","the","scaffold","is","labeled","as","having","density","1","if","the","point","in","the","center","of","it","evaluates","to","greater","than","`","scaffold_empty_space_threshold","`",".","3D","max","pooling","is","performed","on","the","densities","of","the","points","in","3D",".","Scaffold","features","are","off","by","default",".","Members",":","voxel_grid_density","(","VoxelGridBase",")",":","voxel","grid","to","use","for","density","estimation","voxel_grid_color","(","VoxelGridBase",")",":","voxel","grid","to","use","for","color","estimation","harmonic_embedder_xyz_density","(","HarmonicEmbedder",")",":","Function","to","transform","the","outputs","of","the","voxel_grid_density","harmonic_embedder_xyz_color","(","HarmonicEmbedder",")",":","Function","to","transform","the","outputs","of","the","voxel_grid_color","for","density","harmonic_embedder_dir_color","(","HarmonicEmbedder",")",":","Function","to","transform","the","outputs","of","the","voxel_grid_color","for","color","decoder_density","(","DecoderFunctionBase",")",":","decoder","function","to","use","for","density","estimation","color_density","(","DecoderFunctionBase",")",":","decoder","function","to","use","for","color","estimation","use_multiple_streams","(","bool",")",":","if","you","want","the","density","and","color","calculations","to","run","on","different","cuda","streams","set","this","to","True",".","Default","True",".","xyz_ray_dir_in_camera_coords","(","bool",")",":","This","is","true","if","the","directions","are","given","in","camera","coordinates",".","Default","False",".","voxel_grid_scaffold","(","VoxelGridModule",")",":","which","holds","the","scaffold",".","Extents","and","translation","of","it","are","set","to","those","of","voxel_grid_density",".","scaffold_calculating_epochs","(","Tuple","[","int",",","...","]",")",":","at","which","epochs","to","recalculate","the","scaffold",".","(","The","scaffold","will","be","created","automatically","at","the","beginning","of","the","calculation",".",")","scaffold_resolution","(","Tuple","[","int",",","int",",","int","]",")",":","(","width",",","height",",","depth",")","of","the","underlying","voxel","grid","which","stores","scaffold","scaffold_empty_space_threshold","(","float",")",":","if","`","self._get_density","`","evaluates","to","less","than","this","it","will","be","considered","as","empty","space","and","the","scaffold","at","that","point","would","evaluate","as","empty","space",".","scaffold_occupancy_chunk_size","(","str","or","int",")",":","Number","of","xy","scaffold","planes","to","calculate","at","the","same","time",".","To","calculate","the","scaffold","we","need","to","query","`","_get_density","(",")","`","at","every","voxel",",","this","calculation","can","be","split","into","scaffold","depth","number","of","xy","plane","calculations","if","you","want","the","lowest","memory","usage",",","one","calculation","to","calculate","the","whole","scaffold",",","but","with","higher","memory","footprint","or","any","other","number","of","planes",".","Setting","to","a","non-positive","number","calculates","all","planes","at","the","same","time",".","Defaults","to","'-1","'","(","=calculating","all","planes",")",".","scaffold_max_pool_kernel_size","(","int",")",":","Size","of","the","pooling","region","to","use","when","calculating","the","scaffold",".","Defaults","to","3.","scaffold_filter_points","(","bool",")",":","If","set","to","True","the","points","will","be","filtered","using","`","self.voxel_grid_scaffold","`",".","Filtered","points","will","be","predicted","as","having","0","density","and","(","0",",","0",",","0",")","color",".","The","points","which","were","not","evaluated","as","empty","space","will","be","passed","through","the","steps","outlined","above",".","volume_cropping_epochs",":","on","which","epochs","to","crop","the","voxel","grids","to","fit","the","object's","bounding","box",".","Scaffold","has","to","be","calculated","before","cropping","."]
docstr: The forward function accepts the parametrizations of 3D points sampled along
  projection rays. The forward pass is responsible for attaching a 3D vector
  and a 1D scalar representing the point's RGB color and opacity respectively.
  
  Args:
      ray_bundle: An ImplicitronRayBundle object containing the following variables:
          origins: A tensor of shape `(minibatch, ..., 3)` denoting the
              origins of the sampling rays in world coords.
          directions: A tensor of shape `(minibatch, ..., 3)`
              containing the direction vectors of sampling rays in world coords.
          lengths: A tensor of shape `(minibatch, ..., num_points_per_ray)`
              containing the lengths at which the rays are sampled.
      fun_viewpool: an optional callback with the signature
              fun_fiewpool(points) -> pooled_features
          where points is a [N_TGT x N x 3] tensor of world coords,
          and pooled_features is a [N_TGT x ... x N_SRC x latent_dim] tensor
          of the features pooled from the context images.
      camera: A camera model which will be used to transform the viewing
          directions
  
  Returns:
      rays_densities: A tensor of shape `(minibatch, ..., num_points_per_ray, 1)`
          denoting the opacitiy of each ray point.
      rays_colors: A tensor of shape `(minibatch, ..., num_points_per_ray, 3)`
          denoting the color of each ray point.
docstr_tok: ["The","forward","function","accepts","the","parametrizations","of","3D","points","sampled","along","projection","rays",".","The","forward","pass","is","responsible","for","attaching","a","3D","vector","and","a","1D","scalar","representing","the","point","'s","RGB","color","and","opacity","respectively",".","Args",":","ray_bundle",":","An","ImplicitronRayBundle","object","containing","the","following","variables",":","origins",":","A","tensor","of","shape","`","(","minibatch",",","...",",","3",")","`","denoting","the","origins","of","the","sampling","rays","in","world","coords",".","directions",":","A","tensor","of","shape","`","(","minibatch",",","...",",","3",")","`","containing","the","direction","vectors","of","sampling","rays","in","world","coords",".","lengths",":","A","tensor","of","shape","`","(","minibatch",",","...",",","num_points_per_ray",")","`","containing","the","lengths","at","which","the","rays","are","sampled",".","fun_viewpool",":","an","optional","callback","with","the","signature","fun_fiewpool","(","points",")","-",">","pooled_features","where","points","is","a","[","N_TGT","x","N","x","3","]","tensor","of","world","coords",",","and","pooled_features","is","a","[","N_TGT","x","...","x","N_SRC","x","latent_dim","]","tensor","of","the","features","pooled","from","the","context","images",".","camera",":","A","camera","model","which","will","be","used","to","transform","the","viewing","directions","Returns",":","rays_densities",":","A","tensor","of","shape","`","(","minibatch",",","...",",","num_points_per_ray",",","1",")","`","denoting","the","opacitiy","of","each","ray","point",".","rays_colors",":","A","tensor","of","shape","`","(","minibatch",",","...",",","num_points_per_ray",",","3",")","`","denoting","the","color","of","each","ray","point","."]
code: def forward(
      self,
      ray_bundle: ImplicitronRayBundle,
      fun_viewpool=None,
      camera: Optional[CamerasBase] = None,
      global_code=None,
      **kwargs,
  ) -> Tuple[torch.Tensor, torch.Tensor, Dict]:
      """
      The forward function accepts the parametrizations of 3D points sampled along
      projection rays. The forward pass is responsible for attaching a 3D vector
      and a 1D scalar representing the point's RGB color and opacity respectively.
  
      Args:
          ray_bundle: An ImplicitronRayBundle object containing the following variables:
              origins: A tensor of shape `(minibatch, ..., 3)` denoting the
                  origins of the sampling rays in world coords.
              directions: A tensor of shape `(minibatch, ..., 3)`
                  containing the direction vectors of sampling rays in world coords.
              lengths: A tensor of shape `(minibatch, ..., num_points_per_ray)`
                  containing the lengths at which the rays are sampled.
          fun_viewpool: an optional callback with the signature
                  fun_fiewpool(points) -> pooled_features
              where points is a [N_TGT x N x 3] tensor of world coords,
              and pooled_features is a [N_TGT x ... x N_SRC x latent_dim] tensor
              of the features pooled from the context images.
          camera: A camera model which will be used to transform the viewing
              directions
  
      Returns:
          rays_densities: A tensor of shape `(minibatch, ..., num_points_per_ray, 1)`
              denoting the opacitiy of each ray point.
          rays_colors: A tensor of shape `(minibatch, ..., num_points_per_ray, 3)`
              denoting the color of each ray point.
      """
      # ########## convert the ray parametrizations to world coordinates ########## #
      # points.shape = [minibatch x n_rays_width x n_rays_height x pts_per_ray x 3]
      # pyre-ignore[6]
      points = ray_bundle_to_ray_points(ray_bundle)
      directions = ray_bundle.directions.reshape(-1, 3)
      input_shape = points.shape
      num_points_per_ray = input_shape[-2]
      points = points.view(-1, 3)
      non_empty_points = None
  
      # ########## filter the points using the scaffold ########## #
      if self._scaffold_ready and self.scaffold_filter_points:
          with torch.no_grad():
              non_empty_points = (
                  self.voxel_grid_scaffold(points)[..., 0] > 0
              )
          points = points[non_empty_points]
          if len(points) == 0:
              warnings.warn(
                  "The scaffold has filtered all the points."
                  "The voxel grids and decoding functions will not be run."
              )
              return (
                  points.new_zeros((*input_shape[:-1], 1)),
                  points.new_zeros((*input_shape[:-1], 3)),
                  {},
              )
  
      # ########## calculate color and density ########## #
      rays_densities, rays_colors = self._calculate_density_and_color(
          points,
          directions,
          camera,
          non_empty_points,
          num_points_per_ray,
      )
  
      if not (self._scaffold_ready and self.scaffold_filter_points):
          return (
              rays_densities.view(
                  (*input_shape[:-1], rays_densities.shape[-1])
              ),
              rays_colors.view(
                  (*input_shape[:-1], rays_colors.shape[-1])
              ),
              {},
          )
  
      # ########## merge scaffold calculated points ########## #
      # Create a zeroed tensor corresponding to a point with density=0 and fill it
      # with calculated density for points which are not in empty space. Do the
      # same for color
      rays_densities_combined = rays_densities.new_zeros(
          (math.prod(input_shape[:-1]), rays_densities.shape[-1])
      )
      rays_colors_combined = rays_colors.new_zeros(
          (math.prod(input_shape[:-1]), rays_colors.shape[-1])
      )
      assert non_empty_points is not None
      rays_densities_combined[non_empty_points] = rays_densities
      rays_colors_combined[non_empty_points] = rays_colors
  
      return (
          rays_densities_combined.view(
              (*input_shape[:-1], rays_densities.shape[-1])
          ),
          rays_colors_combined.view(
              (*input_shape[:-1], rays_colors.shape[-1])
          ),
          {},
      )
  
imports: import logging
  import math
  import warnings
  from dataclasses import fields
  from typing import Callable, Dict, Optional, Tuple
  import torch
  from omegaconf import DictConfig
  from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase
  from pytorch3d.implicitron.models.implicit_function.decoding_functions import DecoderFunctionBase
  from pytorch3d.implicitron.models.implicit_function.voxel_grid import VoxelGridModule
  from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
  from pytorch3d.implicitron.tools.config import enable_get_default_args, get_default_args_field, registry, run_auto_creation
  from pytorch3d.renderer import ray_bundle_to_ray_points
  from pytorch3d.renderer.cameras import CamerasBase
  from pytorch3d.renderer.implicit import HarmonicEmbedding
label_desc: Use image node_id 2 for calling the VoxelGridImplicitFunction obj's underlying member method code with example usage: obj.forward(ray_bundle, fun_viewpool, camera, global_code) and returns: rays_densities_combined,  rays_colors_combined,  dict, rays_densities,  rays_colors,  dict, points,  points,  dict