class_docstr: Nerf NN Model with both coarse and fine MLPs.
  
  Attributes:
    num_coarse_samples: int, the number of samples for coarse nerf.
    num_fine_samples: int, the number of samples for fine nerf.
    use_viewdirs: bool, use viewdirs as a condition.
    near: float, near clip.
    far: float, far clip.
    noise_std: float, std dev of noise added to regularize sigma output.
    nerf_trunk_depth: int, the depth of the first part of MLP.
    nerf_trunk_width: int, the width of the first part of MLP.
    nerf_rgb_branch_depth: int, the depth of the second part of MLP.
    nerf_rgb_branch_width: int, the width of the second part of MLP.
    activation: the activation function used in the MLP.
    sigma_activation: the activation function applied to the sigma density.
    nerf_skips: which layers to add skip layers in the NeRF model.
    alpha_channels: int, the number of alpha_channelss.
    rgb_channels: int, the number of rgb_channelss.
    use_stratified_sampling: use stratified sampling.
    use_white_background: composite rendering on to a white background.
    num_nerf_point_freqs: degree of positional encoding for positions.
    num_nerf_viewdir_freqs: degree of positional encoding for viewdirs.
    use_linear_disparity: sample linearly in disparity rather than depth.
    num_appearance_embeddings: the number of appearance exemplars.
    num_appearance_features: the dimension size for the appearance code.
    num_camera_embeddings: the number of camera exemplars.
    num_camera_features: the dimension size for the camera code.
    num_warp_freqs: the number of frequencies for the warp positional encoding.
    num_warp_embeddings: the number of embeddings for the warp GLO encoder.
    num_warp_features: the number of features for the warp GLO encoder.
    use_appearance_metadata: if True use the appearance metadata.
    use_camera_metadata: if True use the camera metadata.
    use_warp: whether to use the warp field or not.
    use_warp_jacobian: if True the model computes and returns the Jacobian of
      the warped points.
    use_weights: if True return the density weights.
    use_trunk_condition: whether to feed the appearance metadata to the trunk.
    use_alpha_condition: whether to feed the appearance metadata to the alpha
      branch.
    use_rgb_condition: whether to feed the appearance metadata to the rgb
      branch.
    warp_kwargs: extra keyword arguments for the warp field.
class_docstr_tok: ["Nerf","NN","Model","with","both","coarse","and","fine","MLPs",".","Attributes",":","num_coarse_samples",":","int",",","the","number","of","samples","for","coarse","nerf",".","num_fine_samples",":","int",",","the","number","of","samples","for","fine","nerf",".","use_viewdirs",":","bool",",","use","viewdirs","as","a","condition",".","near",":","float",",","near","clip",".","far",":","float",",","far","clip",".","noise_std",":","float",",","std","dev","of","noise","added","to","regularize","sigma","output",".","nerf_trunk_depth",":","int",",","the","depth","of","the","first","part","of","MLP",".","nerf_trunk_width",":","int",",","the","width","of","the","first","part","of","MLP",".","nerf_rgb_branch_depth",":","int",",","the","depth","of","the","second","part","of","MLP",".","nerf_rgb_branch_width",":","int",",","the","width","of","the","second","part","of","MLP",".","activation",":","the","activation","function","used","in","the","MLP",".","sigma_activation",":","the","activation","function","applied","to","the","sigma","density",".","nerf_skips",":","which","layers","to","add","skip","layers","in","the","NeRF","model",".","alpha_channels",":","int",",","the","number","of","alpha_channelss",".","rgb_channels",":","int",",","the","number","of","rgb_channelss",".","use_stratified_sampling",":","use","stratified","sampling",".","use_white_background",":","composite","rendering","on","to","a","white","background",".","num_nerf_point_freqs",":","degree","of","positional","encoding","for","positions",".","num_nerf_viewdir_freqs",":","degree","of","positional","encoding","for","viewdirs",".","use_linear_disparity",":","sample","linearly","in","disparity","rather","than","depth",".","num_appearance_embeddings",":","the","number","of","appearance","exemplars",".","num_appearance_features",":","the","dimension","size","for","the","appearance","code",".","num_camera_embeddings",":","the","number","of","camera","exemplars",".","num_camera_features",":","the","dimension","size","for","the","camera","code",".","num_warp_freqs",":","the","number","of","frequencies","for","the","warp","positional","encoding",".","num_warp_embeddings",":","the","number","of","embeddings","for","the","warp","GLO","encoder",".","num_warp_features",":","the","number","of","features","for","the","warp","GLO","encoder",".","use_appearance_metadata",":","if","True","use","the","appearance","metadata",".","use_camera_metadata",":","if","True","use","the","camera","metadata",".","use_warp",":","whether","to","use","the","warp","field","or","not",".","use_warp_jacobian",":","if","True","the","model","computes","and","returns","the","Jacobian","of","the","warped","points",".","use_weights",":","if","True","return","the","density","weights",".","use_trunk_condition",":","whether","to","feed","the","appearance","metadata","to","the","trunk",".","use_alpha_condition",":","whether","to","feed","the","appearance","metadata","to","the","alpha","branch",".","use_rgb_condition",":","whether","to","feed","the","appearance","metadata","to","the","rgb","branch",".","warp_kwargs",":","extra","keyword","arguments","for","the","warp","field","."]
code: def setup(self):
      if self.use_warp:
          self.warp_field = self.create_warp_field(
              self, num_batch_dims=2
          )
  
      self.point_encoder = model_utils.vmap_module(
          modules.SinusoidalEncoder, num_batch_dims=2
      )(num_freqs=self.num_nerf_point_freqs)
      self.viewdir_encoder = model_utils.vmap_module(
          modules.SinusoidalEncoder, num_batch_dims=1
      )(num_freqs=self.num_nerf_viewdir_freqs)
      if self.use_appearance_metadata:
          self.appearance_encoder = glo.GloEncoder(
              num_embeddings=self.num_appearance_embeddings,
              features=self.num_appearance_features,
          )
      if self.use_camera_metadata:
          self.camera_encoder = glo.GloEncoder(
              num_embeddings=self.num_camera_embeddings,
              features=self.num_camera_features,
          )
  
      nerf_mlps = {
          "coarse": modules.NerfMLP(
              trunk_depth=self.nerf_trunk_depth,
              trunk_width=self.nerf_trunk_width,
              rgb_branch_depth=self.nerf_rgb_branch_depth,
              rgb_branch_width=self.nerf_rgb_branch_width,
              activation=self.activation,
              skips=self.nerf_skips,
              alpha_channels=self.alpha_channels,
              rgb_channels=self.rgb_channels,
          )
      }
      if self.num_fine_samples > 0:
          nerf_mlps["fine"] = modules.NerfMLP(
              trunk_depth=self.nerf_trunk_depth,
              trunk_width=self.nerf_trunk_width,
              rgb_branch_depth=self.nerf_rgb_branch_depth,
              rgb_branch_width=self.nerf_rgb_branch_width,
              activation=self.activation,
              skips=self.nerf_skips,
              alpha_channels=self.alpha_channels,
              rgb_channels=self.rgb_channels,
          )
      self.nerf_mlps = nerf_mlps
  
imports: from typing import Any, Dict, Mapping, Optional, Tuple, Sequence
  from flax import linen
  import immutabledict
  from jax import random
  import jax.numpy
  from nerfies import configs
  from nerfies import glo
  from nerfies import model_utils
  from nerfies import modules
  from nerfies import types
  from nerfies import warping
label_desc: Use image node_id 5 for calling the NerfModel obj's underlying member method code with example usage: obj.setup() without return types