class_docstr: Implements the sin-cos positional encoding.
class_docstr_tok: ["Implements","the","sin-cos","positional","encoding","."]
docstr: Returns the input activations, with added positional information.
docstr_tok: ["Returns","the","input","activations",",","with","added","positional","information","."]
code: 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
  
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 3 for calling the SinCosPositionalEncoding obj's underlying member method code with example usage: obj.forward(inputs) and returns: unknown, unknown, unknown