(self, input_signature)
Randomly initializes this layer's weights.
Randomly initializes this layer's weights.
Randomly initializes this layer s weights .
def init_weights_and_state(self, input_signature):
    """Randomly initializes this layer's weights."""
    d_model = input_signature[-1].shape[-1]
    shape_w1 = (d_model, self._d_ff)
    shape_w2 = (self._d_ff, d_model)
    shape_b2 = (d_model,)

    rng_w1, rng_w2, rng_b2 = fastmath.random.split(self.rng, 3)
    if base.N_WEIGHTS_SHARDS > 1:
      # In sharded-weights mode, put the weights on CPU on init
      # as they will be sharded later.
      w1 = tl.on_cpu(self._kernel_initializer(shape_w1, rng_w1))
      w2 = tl.on_cpu(self._kernel_initializer(shape_w2, rng_w2))
    else:
      w1 = self._kernel_initializer(shape_w1, rng_w1)
      w2 = self._kernel_initializer(shape_w2, rng_w2)

    b2 = self._bias_initializer(shape_b2, rng_b2)
    if self._use_bfloat16:
      b2 = b2.astype(jnp.bfloat16)
    if self._use_bfloat16 or self._big_weights_in_bfloat16:
      w1 = w1.astype(jnp.bfloat16)
      w2 = w2.astype(jnp.bfloat16)

    w1 = jnp.reshape(w1, (-1, self._d1, self._d2))
    w2 = jnp.reshape(w2, (self._d2, self._d1, -1))

    if self._kernel_scaling:
      # This keeps expected variance of the output regardless of N.
      w2 = w2 * (self._n_elements_in_block ** 0.5)

    self.weights = (w1, w2, b2)
def init_weights_and_state ( self , input_signature ) : d_model = input_signature [ - 1 ] . shape [ - 1 ] shape_w1 = ( d_model , self . _d_ff ) shape_w2 = ( self . _d_ff , d_model ) shape_b2 = ( d_model , ) rng_w1 , rng_w2 , rng_b2 = fastmath . random . split ( self . rng , 3 ) if base . N_WEIGHTS_SHARDS > 1 : # In sharded-weights mode, put the weights on CPU on init # as they will be sharded later. w1 = tl . on_cpu ( self . _kernel_initializer ( shape_w1 , rng_w1 ) ) w2 = tl . on_cpu ( self . _kernel_initializer ( shape_w2 , rng_w2 ) ) else : w1 = self . _kernel_initializer ( shape_w1 , rng_w1 ) w2 = self . _kernel_initializer ( shape_w2 , rng_w2 ) b2 = self . _bias_initializer ( shape_b2 , rng_b2 ) if self . _use_bfloat16 : b2 = b2 . astype ( jnp . bfloat16 ) if self . _use_bfloat16 or self . _big_weights_in_bfloat16 : w1 = w1 . astype ( jnp . bfloat16 ) w2 = w2 . astype ( jnp . bfloat16 ) w1 = jnp . reshape ( w1 , ( - 1 , self . _d1 , self . _d2 ) ) w2 = jnp . reshape ( w2 , ( self . _d2 , self . _d1 , - 1 ) ) if self . _kernel_scaling : # This keeps expected variance of the output regardless of N. w2 = w2 * ( self . _n_elements_in_block ** 0.5 ) self . weights = ( w1 , w2 , b2 )