(self, inputs)
Returns the input activations, with added positional information.
Returns the input activations, with added positional information.
Returns the input activations with added positional information .
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
def forward ( self , inputs ) : 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