def forward(self, inputs):
    """Returns the input activations, with added positional information."""
    weights = self.weights
    if self._d_feature is not None:
      weights, ff = weights
      weights = jnp.dot(weights[:inputs.shape[1], :], ff)
    if len(weights.shape) < 3:  # old checkpoints have 1 in first dim already
      weights = weights[None, :, :]  # [1, self._max_len, d_feature]
    if self._mode != 'predict':
      x = inputs
      symbol_size = jnp.shape(x)[1]
      if self._mode != 'train' or self._start_from_zero_prob >= 1.0:
        px = weights[:, :symbol_size, :]
      else:
        rng1, rng2 = fastmath.random.split(self.rng, 2)
        start = fastmath.random.randint(rng1, (), 0, self._max_offset_to_add)
        start_from_zero = fastmath.random.uniform(rng2, (), jnp.float32, 0, 1)
        start = jnp.where(start_from_zero < self._start_from_zero_prob,
                          jnp.zeros((), dtype=jnp.int32), start)
        px = fastmath.dynamic_slice_in_dim(weights, start, symbol_size,
                                           axis=1)
      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 stores the index of the current
      # position and increments it on each call.
      emb = fastmath.dynamic_slice_in_dim(
          weights, self.state, inputs.shape[1], axis=1)
      self.state += inputs.shape[1]
      return inputs + emb
Returns the input activations, with added positional information.
trax/layers/attention.py