class_docstr: Network that predicts warps as a translation field.
  
  References:
    https://en.wikipedia.org/wiki/Vector_potential
    https://en.wikipedia.org/wiki/Helmholtz_decomposition
  
  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 output 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 output layer.
class_docstr_tok: ["Network","that","predicts","warps","as","a","translation","field",".","References",":","https",":","\/\/en.wikipedia.org\/wiki\/Vector_potential","https",":","\/\/en.wikipedia.org\/wiki\/Helmholtz_decomposition","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","output","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","output","layer","."]
docstr: Warp the given points using a warp field.
  
  Args:
    points: the points to warp.
    metadata: metadata indices if metadata_encoded is False else pre-encoded
      metadata.
    extra: extra parameters used in the warp field e.g., the warp alpha.
    return_jacobian: if True compute and return the Jacobian of the warp.
    metadata_encoded: if True assumes the metadata is already encoded.
  
  Returns:
    The warped points and the Jacobian of the warp if `return_jacobian` is
      True.
docstr_tok: ["Warp","the","given","points","using","a","warp","field",".","Args",":","points",":","the","points","to","warp",".","metadata",":","metadata","indices","if","metadata_encoded","is","False","else","pre-encoded","metadata",".","extra",":","extra","parameters","used","in","the","warp","field","e.g.",",","the","warp","alpha",".","return_jacobian",":","if","True","compute","and","return","the","Jacobian","of","the","warp",".","metadata_encoded",":","if","True","assumes","the","metadata","is","already","encoded",".","Returns",":","The","warped","points","and","the","Jacobian","of","the","warp","if","`","return_jacobian","`","is","True","."]
code: def __call__(
      self,
      points: jnp.ndarray,
      metadata: jnp.ndarray,
      extra: Dict[str, Any],
      return_jacobian: bool = False,
      metadata_encoded: bool = False,
  ):
      """Warp the given points using a warp field.
  
      Args:
        points: the points to warp.
        metadata: metadata indices if metadata_encoded is False else pre-encoded
          metadata.
        extra: extra parameters used in the warp field e.g., the warp alpha.
        return_jacobian: if True compute and return the Jacobian of the warp.
        metadata_encoded: if True assumes the metadata is already encoded.
  
      Returns:
        The warped points and the Jacobian of the warp if `return_jacobian` is
          True.
      """
      if metadata_encoded:
          metadata_embed = metadata
      else:
          metadata_embed = self.encode_metadata(
              metadata, extra.get("time_alpha")
          )
  
      out = {"warped_points": self.warp(points, metadata_embed, extra)}
  
      if return_jacobian:
          jac_fn = jax.jacfwd(
              lambda *x: self.warp(*x)[..., :3], argnums=0
          )
          out["jacobian"] = jac_fn(points, metadata_embed, extra)
  
      return out
  
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 4 for calling the TranslationField obj's underlying member method code with example usage: obj.__call__(points, metadata, extra, return_jacobian, metadata_encoded) and returns: out