class_docstr: Network that predicts warps as an SE(3) field.
  
  Attributes:
    points_encoder: the positional encoder for the points.
    metadata_encoder: an encoder for metadata.
    alpha: the alpha for the positional encoding.
    skips: the index of the layers with skip connections.
    depth: the depth of the network excluding the logit layer.
    hidden_channels: the width of the network hidden layers.
    activation: the activation for each layer.
    metadata_encoded: whether the metadata parameter is pre-encoded or not.
    hidden_initializer: the initializer for the hidden layers.
    output_initializer: the initializer for the last logit layer.
class_docstr_tok: ["Network","that","predicts","warps","as","an","SE","(","3",")","field",".","Attributes",":","points_encoder",":","the","positional","encoder","for","the","points",".","metadata_encoder",":","an","encoder","for","metadata",".","alpha",":","the","alpha","for","the","positional","encoding",".","skips",":","the","index","of","the","layers","with","skip","connections",".","depth",":","the","depth","of","the","network","excluding","the","logit","layer",".","hidden_channels",":","the","width","of","the","network","hidden","layers",".","activation",":","the","activation","for","each","layer",".","metadata_encoded",":","whether","the","metadata","parameter","is","pre-encoded","or","not",".","hidden_initializer",":","the","initializer","for","the","hidden","layers",".","output_initializer",":","the","initializer","for","the","last","logit","layer","."]
code: def warp(
      self,
      points: jnp.ndarray,
      metadata_embed: jnp.ndarray,
      extra: Dict[str, Any],
  ):
      points_embed = self.points_encoder(
          points, alpha=extra.get("alpha")
      )
      inputs = jnp.concatenate([points_embed, metadata_embed], axis=-1)
      trunk_output = self.trunk(inputs)
  
      w = self.branches["w"](trunk_output)
      v = self.branches["v"](trunk_output)
      theta = jnp.linalg.norm(w, axis=-1)
      w = w / theta[..., jnp.newaxis]
      v = v / theta[..., jnp.newaxis]
      screw_axis = jnp.concatenate([w, v], axis=-1)
      transform = rigid.exp_se3(screw_axis, theta)
  
      warped_points = points
      if self.use_pivot:
          pivot = self.branches["p"](trunk_output)
          warped_points = warped_points + pivot
  
      warped_points = rigid.from_homogenous(
          transform @ rigid.to_homogenous(warped_points)
      )
  
      if self.use_pivot:
          warped_points = warped_points - pivot
  
      if self.use_translation:
          t = self.branches["t"](trunk_output)
          warped_points = warped_points + t
  
      return warped_points
  
imports: from typing import Any, Iterable, Optional, Dict
  from flax import linen
  import jax
  import jax.numpy
  from nerfies import glo
  from nerfies import model_utils
  from nerfies import modules
  from nerfies import rigid_body
  from nerfies import types
label_desc: Use image node_id 3 for calling the SE3Field obj's underlying member method code with example usage: obj.warp(points, metadata_embed, extra) and returns: warped_points