(self, x)
return jnp.reshape(res, x_shape)
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.
    """
    m1, w1, w2, b2 = self.weights
    x_shape = x.shape
    x = jnp.reshape(x, [-1, x_shape[-1]])  # Easier to operate on flattened x.

    # Q: check if we need bias and/or put relu after the m1 dot?
    mask_logits = jnp.dot(x, m1)
    # Softmax.
    mask_logsumexp = fastmath.logsumexp(mask_logits, axis=-1, keepdims=True)
    log_mask = mask_logits - mask_logsumexp
    mask = jnp.exp(log_mask)
    # Gumbel noise to allow sampling from the softmax.
    rng1, _ = fastmath.random.split(self.rng, 2)
    u = fastmath.random.uniform(rng1, mask.shape, jnp.float32, 1e-6, 1.0 - 1e-6)
    g = -jnp.log(-jnp.log(u))
    selected_experts = jnp.argmax(log_mask + g * self._temperature, axis=-1)
    quant_mask = tl.one_hot(selected_experts, self._n_experts)
    quant_mask = fastmath.stop_gradient(quant_mask)
    quant_mask *= mask  # go to just the selected expert
    quant_mask = jnp.reshape(quant_mask, [-1, self._n_experts, 1])
    batch_size = quant_mask.shape[0]

    if self._mode == 'predict' and batch_size == 1:
      mask_flat = jnp.reshape(mask, [-1, self._n_experts])
      selected_flat = jnp.reshape(selected_experts, [-1])
      selected_mask_flat = mask_flat[np.arange(selected_flat.size),
                                     selected_flat]
      # This implementation mimicks inference for batch_size 1.
      start_idx = selected_experts[0] * self._n_elements_in_block
      # w1 is [d_model, d_ff], w is [d_model, n_elements_in_block]
      w = fastmath.dynamic_slice(w1, [0, start_idx],
                                 [w1.shape[0], self._n_elements_in_block])
      mid = jnp.dot(x, w)
      mid *= jnp.reshape(selected_mask_flat, mid.shape[:-1])[..., None]
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      # w2 is [d_ff, d_model], v is [n_elements_in_block, d_model]
      v = fastmath.dynamic_slice(w2, [start_idx, 0],
                                 [self._n_elements_in_block, w2.shape[-1]])
      v = jnp.reshape(v, [self._n_elements_in_block, -1])
      res = jnp.dot(relu, v) + b2
    else:
      expanded_mask = jnp.broadcast_to(
          quant_mask,
          (quant_mask.shape[0], quant_mask.shape[1], self._n_elements_in_block))
      expanded_mask = jnp.reshape(expanded_mask, (-1, self._d_ff))
      mid = jnp.dot(x, w1) * expanded_mask  # [joint_batch, d_ff]
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      res = jnp.dot(relu, w2) + b2

    return jnp.reshape(res, x_shape)
def forward ( self , x ) : m1 , w1 , w2 , b2 = self . weights x_shape = x . shape x = jnp . reshape ( x , [ - 1 , x_shape [ - 1 ] ] ) # Easier to operate on flattened x. # Q: check if we need bias and/or put relu after the m1 dot? mask_logits = jnp . dot ( x , m1 ) # Softmax. mask_logsumexp = fastmath . logsumexp ( mask_logits , axis = - 1 , keepdims = True ) log_mask = mask_logits - mask_logsumexp mask = jnp . exp ( log_mask ) # Gumbel noise to allow sampling from the softmax. rng1 , _ = fastmath . random . split ( self . rng , 2 ) u = fastmath . random . uniform ( rng1 , mask . shape , jnp . float32 , 1e-6 , 1.0 - 1e-6 ) g = - jnp . log ( - jnp . log ( u ) ) selected_experts = jnp . argmax ( log_mask + g * self . _temperature , axis = - 1 ) quant_mask = tl . one_hot ( selected_experts , self . _n_experts ) quant_mask = fastmath . stop_gradient ( quant_mask ) quant_mask *= mask # go to just the selected expert quant_mask = jnp . reshape ( quant_mask , [ - 1 , self . _n_experts , 1 ] ) batch_size = quant_mask . shape [ 0 ] if self . _mode == 'predict' and batch_size == 1 : mask_flat = jnp . reshape ( mask , [ - 1 , self . _n_experts ] ) selected_flat = jnp . reshape ( selected_experts , [ - 1 ] ) selected_mask_flat = mask_flat [ np . arange ( selected_flat . size ) , selected_flat ] # This implementation mimicks inference for batch_size 1. start_idx = selected_experts [ 0 ] * self . _n_elements_in_block # w1 is [d_model, d_ff], w is [d_model, n_elements_in_block] w = fastmath . dynamic_slice ( w1 , [ 0 , start_idx ] , [ w1 . shape [ 0 ] , self . _n_elements_in_block ] ) mid = jnp . dot ( x , w ) mid *= jnp . reshape ( selected_mask_flat , mid . shape [ : - 1 ] ) [ ... , None ] relu = jnp . where ( mid <= 0 , jnp . zeros_like ( mid ) , mid ) # w2 is [d_ff, d_model], v is [n_elements_in_block, d_model] v = fastmath . dynamic_slice ( w2 , [ start_idx , 0 ] , [ self . _n_elements_in_block , w2 . shape [ - 1 ] ] ) v = jnp . reshape ( v , [ self . _n_elements_in_block , - 1 ] ) res = jnp . dot ( relu , v ) + b2 else : expanded_mask = jnp . broadcast_to ( quant_mask , ( quant_mask . shape [ 0 ] , quant_mask . shape [ 1 ] , self . _n_elements_in_block ) ) expanded_mask = jnp . reshape ( expanded_mask , ( - 1 , self . _d_ff ) ) mid = jnp . dot ( x , w1 ) * expanded_mask # [joint_batch, d_ff] relu = jnp . where ( mid <= 0 , jnp . zeros_like ( mid ) , mid ) res = jnp . dot ( relu , w2 ) + b2 return jnp . reshape ( res , x_shape )