def forward(self, inputs):
    """Returns the input activations, with added positional information."""
    if self._mode != 'predict':
      x = inputs
      length = jnp.shape(x)[1]
      if self._mode != 'train':
        start = 0
      else:
        rng1, rng2 = fastmath.random.split(self.rng, 2)
        start = fastmath.random.randint(rng1, (), 0, self._add_offset)
        start_from_nonzero = fastmath.random.randint(
            rng2, (), 0, self._start_from_zero_one_in)
        start_from_nonzero = jnp.minimum(1, start_from_nonzero)
        start *= start_from_nonzero
      px = self._sincos(start, length, inputs.shape[2])
      if self._dropout == 0:
        return x + px
      else:
        noise_shape = list(px.shape)
        for dim in self._dropout_broadcast_dims:
          noise_shape[dim] = 1
        keep_prob = 1.0 - self._dropout
        keep = fastmath.random.bernoulli(self.rng, keep_prob,
                                         tuple(noise_shape))
        multiplier = keep.astype(x.dtype) / keep_prob
        return x + px * multiplier
    else:
      if self._dropout != 0:
        raise ValueError(f'In predict mode, but dropout rate '
                         f'({self._dropout}) is not zero.')

      # State in this class is only used for fast inference. In that case,
      # the model is called with consecutive elements position-by-position.
      # This positional encoding layer needs to store the index of the current
      # position then and increment it on each call -- that's how state is used
      # and updated below.
      pe = self._sincos(self.state, inputs.shape[1], inputs.shape[2])
      self.state += inputs.shape[1]
      return inputs + pe
Returns the input activations, with added positional information.
trax/layers/research/position_encodings.py