(self, x)
Executes this layer as part of a forward pass through the model.

    Args:
      x: Tensor of same shape and dtype as the input signature used to
          initialize this layer.

    Returns:
      Tensor of same shape and dtype as the input.
Executes this layer as part of a forward pass through the model.
Executes this layer as part of a forward pass through the model .
def forward(self, x):
    """Executes this layer as part of a forward pass through the model.

    Args:
      x: Tensor of same shape and dtype as the input signature used to
          initialize this layer.

    Returns:
      Tensor of same shape and dtype as the input.
    """
    x, recalled_quant_mask = x
    m1, m2, mb = self.weights

    x_shape = x.shape
    x = jnp.reshape(x, [-1, x_shape[-1]])  # Easier to operate on flattened x.

    # Q: should we add bias and/or put relu after the low-rank m1 dot?
    # Replacing multiplication and reshape by this einsum brings training speed
    # improvement (see also reshape in initialization).
    mask_logits = jnp.einsum('bd,dl,lxy->bxy', x, m1, m2) + mb

    if self._also_return_nondiscrete_output:
      # Softmax.
      mask_logsumexp = fastmath.logsumexp(mask_logits, axis=-1, keepdims=True)
      log_mask = mask_logits - mask_logsumexp
      mask = jnp.exp(log_mask)
      # Gumbel-softmax with straight-through discretization.
      if self._temperature == 0.0:
        quant_mask = jnp.argmax(log_mask, axis=-1)
      else:
        u = fastmath.random.uniform(self.rng, mask.shape, jnp.float32, 1e-6,
                                    1.0 - 1e-6)
        g = -jnp.log(-jnp.log(u))
        quant_mask = jnp.argmax(log_mask + g * self._temperature, axis=-1)
    else:
      quant_mask = jnp.argmax(mask_logits, axis=-1)

    if self._mode == 'train':
      # We use recalled_quant_mask if it's different than -1; otherwise
      # we use a quant_mask which we have just computed.
      quant_mask = jnp.where(recalled_quant_mask == -1,
                             quant_mask, recalled_quant_mask)

    if self._also_return_nondiscrete_output:
      return quant_mask, mask
    else:
      return quant_mask
def forward ( self , x ) : x , recalled_quant_mask = x m1 , m2 , mb = self . weights x_shape = x . shape x = jnp . reshape ( x , [ - 1 , x_shape [ - 1 ] ] ) # Easier to operate on flattened x. # Q: should we add bias and/or put relu after the low-rank m1 dot? # Replacing multiplication and reshape by this einsum brings training speed # improvement (see also reshape in initialization). mask_logits = jnp . einsum ( 'bd,dl,lxy->bxy' , x , m1 , m2 ) + mb if self . _also_return_nondiscrete_output : # Softmax. mask_logsumexp = fastmath . logsumexp ( mask_logits , axis = - 1 , keepdims = True ) log_mask = mask_logits - mask_logsumexp mask = jnp . exp ( log_mask ) # Gumbel-softmax with straight-through discretization. if self . _temperature == 0.0 : quant_mask = jnp . argmax ( log_mask , axis = - 1 ) else : u = fastmath . random . uniform ( self . rng , mask . shape , jnp . float32 , 1e-6 , 1.0 - 1e-6 ) g = - jnp . log ( - jnp . log ( u ) ) quant_mask = jnp . argmax ( log_mask + g * self . _temperature , axis = - 1 ) else : quant_mask = jnp . argmax ( mask_logits , axis = - 1 ) if self . _mode == 'train' : # We use recalled_quant_mask if it's different than -1; otherwise # we use a quant_mask which we have just computed. quant_mask = jnp . where ( recalled_quant_mask == - 1 , quant_mask , recalled_quant_mask ) if self . _also_return_nondiscrete_output : return quant_mask , mask else : return quant_mask