class_docstr: Shape and appearance conditioned network for radiance fields.
class_docstr_tok: ["Shape","and","appearance","conditioned","network","for","radiance","fields","."]
docstr: Main training function for coarse and fine networks.
  
  Args:
    r_org: A tensor of shape `[B, R, 3]` where B is the batch size,
      R is the number of rays and the last dimensios store the ray origin.
    r_dir: A tensor of shape `[B, R, 3]` where B is the batch size,
      R is the number of rays and the last dimensios store the ray direction.
    shape_index: A tensor of shape `[B]` where B is the batch size
    voxels: A tensor of shape `[B, 128, 128, 128, 1]` where B is the batch
      size and the other dimensions contain the voxel grid
    w2v_alpha: A tensor of shape `[B, 3]`
    w2v_beta: A tensor of shape `[B, 3]`
    gt_rgb: A tensor of shape `[B, R, 3]`
  
  Returns:
    A scalar.
docstr_tok: ["Main","training","function","for","coarse","and","fine","networks",".","Args",":","r_org",":","A","tensor","of","shape","`","[","B",",","R",",","3","]","`","where","B","is","the","batch","size",",","R","is","the","number","of","rays","and","the","last","dimensios","store","the","ray","origin",".","r_dir",":","A","tensor","of","shape","`","[","B",",","R",",","3","]","`","where","B","is","the","batch","size",",","R","is","the","number","of","rays","and","the","last","dimensios","store","the","ray","direction",".","shape_index",":","A","tensor","of","shape","`","[","B","]","`","where","B","is","the","batch","size","voxels",":","A","tensor","of","shape","`","[","B",",","128",",","128",",","128",",","1","]","`","where","B","is","the","batch","size","and","the","other","dimensions","contain","the","voxel","grid","w2v_alpha",":","A","tensor","of","shape","`","[","B",",","3","]","`","w2v_beta",":","A","tensor","of","shape","`","[","B",",","3","]","`","gt_rgb",":","A","tensor","of","shape","`","[","B",",","R",",","3","]","`","Returns",":","A","scalar","."]
code: def train_step(
      self,
      r_org,
      r_dir,
      shape_index,
      voxels,
      w2v_alpha,
      w2v_beta,
      gt_rgb,
  ):
      """Main training function for coarse and fine networks.
  
      Args:
        r_org: A tensor of shape `[B, R, 3]` where B is the batch size,
          R is the number of rays and the last dimensios store the ray origin.
        r_dir: A tensor of shape `[B, R, 3]` where B is the batch size,
          R is the number of rays and the last dimensios store the ray direction.
        shape_index: A tensor of shape `[B]` where B is the batch size
        voxels: A tensor of shape `[B, 128, 128, 128, 1]` where B is the batch
          size and the other dimensions contain the voxel grid
        w2v_alpha: A tensor of shape `[B, 3]`
        w2v_beta: A tensor of shape `[B, 3]`
        gt_rgb: A tensor of shape `[B, R, 3]`
  
      Returns:
        A scalar.
      """
      with tf.GradientTape() as tape:
          latent_code = tf.gather(self.latent_code_vars, shape_index)
          rgb_fine = self.inference(
              r_org,
              r_dir,
              latent_code,
              voxels,
              w2v_alpha,
              w2v_beta,
              self.near,
              self.far,
          )
          rgb_fine_loss = utils.l2_loss(rgb_fine, gt_rgb)
          total_loss = rgb_fine_loss
      gradients = tape.gradient(
          total_loss, self.network_vars + [self.latent_code_vars]
      )
      self.optimizer_network.apply_gradients(
          zip(gradients[: len(self.network_vars)], self.network_vars)
      )
      self.optimizer_latent.apply_gradients(
          zip(
              gradients[len(self.network_vars) :],
              [self.latent_code_vars],
          )
      )
      return total_loss
  
imports: from absl import logging
  import numpy
  import tensorflow
  import tensorflow_graphics.geometry.representation.ray
  import tensorflow_graphics.math.feature_representation
  import tensorflow_graphics.projects.radiance_fields.nerf.layers
  import tensorflow_graphics.projects.radiance_fields.sharf.voxel_functions
  import tensorflow_graphics.projects.radiance_fields.utils
  import tensorflow_graphics.rendering.volumetric.ray_radiance
label_desc: Use image node_id 11 for calling the AppearanceNetwork obj's underlying member method code with example usage: obj.train_step(r_org, r_dir, shape_index, voxels, w2v_alpha, w2v_beta, gt_rgb) and returns: total_loss