docstr: Args:
      raymarch_features: Features from the raymarching network of shape
          `(minibatch, ..., self.in_features)`
      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.
  
  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: ["Args",":","raymarch_features",":","Features","from","the","raymarching","network","of","shape","`","(","minibatch",",","...",",","self.in_features",")","`","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",".","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,
      raymarch_features: torch.Tensor,
      ray_bundle: ImplicitronRayBundle,
      camera: Optional[CamerasBase] = None,
      **kwargs,
  ):
      """
      Args:
          raymarch_features: Features from the raymarching network of shape
              `(minibatch, ..., self.in_features)`
          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.
  
      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.
      """
      # raymarch_features.shape = [minibatch x ... x pts_per_ray x 3]
      features = self._net(raymarch_features)
      # features.shape = [minibatch x ... x self.n_hidden_units]
  
      if self.ray_dir_in_camera_coords:
          if camera is None:
              raise ValueError(
                  "Camera must be given if xyz_ray_dir_in_camera_coords"
              )
  
          directions = ray_bundle.directions @ camera.R
      else:
          directions = ray_bundle.directions
  
      # NNs operate on the flattenned rays; reshaping to the correct spatial size
      features = features.reshape(*raymarch_features.shape[:-1], -1)
  
      raw_densities = self._density_layer(features)
  
      rays_colors = self._get_colors(features, directions)
  
      return raw_densities, rays_colors
  
imports: from typing import Any, cast, Optional, Tuple
  import torch
  from omegaconf import DictConfig
  from pytorch3d.common.linear_with_repeat import LinearWithRepeat
  from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
  from pytorch3d.implicitron.third_party import hyperlayers, pytorch_prototyping
  from pytorch3d.implicitron.tools.config import Configurable, 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
  from .base import ImplicitFunctionBase
  from .utils import create_embeddings_for_implicit_function
label_desc: Use image node_id 3 for calling the SRNPixelGenerator obj's underlying member method code with example usage: obj.forward(raymarch_features, ray_bundle, camera) and returns: raw_densities,  rays_colors