class_docstr: Infinite positional encoding.
class_docstr_tok: ["Infinite","positional","encoding","."]
docstr: Initializes the encoding.
  
  The encoding tries to roughly evenly traverse the latent space.
  The recurrence time is dependent on how many bits per dimension you use.
  
  There are two parameters to control randomization:
  - randomizing the origin every 1/drift steps by letting it drift
  - randomizing the origin per call
  
  Args:
    drift: variance in position difference per unit of difference
    affine: whether to randomize the origin every call
    transform: learnable transform after encoding (any/diag/none)
    time_bin_length: Add features AxialPositionalEncoding learns if
      TimeBinCausalAttention is the first layer.
      bin_length should match TBCA.bin_length
      If you set transform='diag', this flag increases your model capacity to
      close to transform='any', though it will still train slower.
    mode: if 'predict', allow evaluating one token at a time
docstr_tok: ["Initializes","the","encoding",".","The","encoding","tries","to","roughly","evenly","traverse","the","latent","space",".","The","recurrence","time","is","dependent","on","how","many","bits","per","dimension","you","use",".","There","are","two","parameters","to","control","randomization",":","-","randomizing","the","origin","every","1\/drift","steps","by","letting","it","drift","-","randomizing","the","origin","per","call","Args",":","drift",":","variance","in","position","difference","per","unit","of","difference","affine",":","whether","to","randomize","the","origin","every","call","transform",":","learnable","transform","after","encoding","(","any\/diag\/none",")","time_bin_length",":","Add","features","AxialPositionalEncoding","learns","if","TimeBinCausalAttention","is","the","first","layer",".","bin_length","should","match","TBCA.bin_length","If","you","set","transform='diag","'",",","this","flag","increases","your","model","capacity","to","close","to","transform='any","'",",","though","it","will","still","train","slower",".","mode",":","if","'predict","'",",","allow","evaluating","one","token","at","a","time"]
code: def __init__(
      self,
      drift=0.03,
      affine=True,
      transform="any",
      time_bin_length=None,
      mode="train",
  ):
      """Initializes the encoding.
  
      The encoding tries to roughly evenly traverse the latent space.
      The recurrence time is dependent on how many bits per dimension you use.
  
      There are two parameters to control randomization:
      - randomizing the origin every 1/drift steps by letting it drift
      - randomizing the origin per call
  
      Args:
        drift: variance in position difference per unit of difference
        affine: whether to randomize the origin every call
        transform: learnable transform after encoding (any/diag/none)
        time_bin_length: Add features AxialPositionalEncoding learns if
          TimeBinCausalAttention is the first layer.
          bin_length should match TBCA.bin_length
          If you set transform='diag', this flag increases your model capacity to
          close to transform='any', though it will still train slower.
        mode: if 'predict', allow evaluating one token at a time
      """
      super().__init__()
      if transform not in ("any", "diag", "none"):
          raise ValueError(transform)
      self._noise_rng = jax.random.split(jax.random.PRNGKey(234234535))[
          0
      ]
      assert self._noise_rng is not None
      self._noise = None
      self._drift = drift
      self._affine = affine
      self._transform = transform
      self._time_bin_length = time_bin_length
      self._mode = mode
  
imports: import logging
  import jax
  import jax.extend
  import numpy
  import trax
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
label_desc: Use image node_id 1 to create a new InfinitePositionalEncoding object from inherited base classes: layer_base with example: obj = InfinitePositionalEncoding(drift, affine, transform, time_bin_length, mode)