class_docstr: Infinite positional encoding.
class_docstr_tok: ["Infinite","positional","encoding","."]
code: def forward(self, inputs):
      rng, state = self.rng, self.state
      d_feature = inputs.shape[-1]
      input_len = inputs.shape[-2]
  
      if self._mode == "predict":
          # Assume all the positions are pretty close to each other.
          index, predict_rng = state
          lo = index.min()
          hi = index.max() + 1
          emb = self._get_embeddings(
              lo=lo, hi=hi, depth=d_feature, rng=predict_rng
          )
          emb = emb[index - lo, jnp.newaxis, :]
          index = index + 1
          state = index, predict_rng
      else:
          emb = self._get_embeddings(
              lo=0, hi=input_len, depth=d_feature, rng=rng
          )
          emb = emb[jnp.newaxis, :input_len, :]
      # TODO(tying): check that XLA swaps matmul(slice(x)) -> slice(matmul(x)),
      # or inline this code into get_embeddings/get_noise
      if self._transform == "diag":
          emb = emb * jax.nn.softplus(self.weights)
      elif self._transform == "any":
          emb = emb @ self.weights
      self.state = state
      return inputs + emb
  
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 4 for calling the InfinitePositionalEncoding obj's underlying member method code with example usage: obj.forward(inputs) and returns: unknown