class_docstr: Implements bare positional encoding.
  
  Positional encoding includes a kind of dropout, if the layer is created in
  ``'train'`` mode with a nonzero ``dropout`` value. For such a layer, on each
  forward pass a subset of sequence positions selected at random will *not*
  receive positional marking.
class_docstr_tok: ["Implements","bare","positional","encoding",".","Positional","encoding","includes","a","kind","of","dropout",",","if","the","layer","is","created","in","``","'train","'","``","mode","with","a","nonzero","``","dropout","``","value",".","For","such","a","layer",",","on","each","forward","pass","a","subset","of","sequence","positions","selected","at","random","will","*","not","*","receive","positional","marking","."]
docstr: Returns the input activations, with added positional information.
code: 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
  
imports: import numpy
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
  from trax.layers.research import sparsity
label_desc: Use image node_id 2 for calling the PositionalEncoding obj's underlying member method code with example usage: obj.forward(inputs) and returns: unknown, unknown, unknown