(self, x)
return original_quant_mask, 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.
    """
    if self._mode == 'train' or self._multiply_by_controller_output:
      quant_mask, mask, x = x
    else:
      quant_mask, x = x
    original_quant_mask = quant_mask

    w1, w2, b2 = self.weights

    if self._mode == 'predict':
      w1 = jnp.transpose(w1, (1, 2, 0))  # dm, d1, d2 -> d1, d2, dm
      w2 = jnp.transpose(w2, (1, 0, 2))  # d2, d1, dm -> d1, d2, dm
    x_shape = x.shape
    x = jnp.reshape(x, [-1, x_shape[-1]])  # Easier to operate on flattened x.

    if self._mode == 'train':
      # Tricks from Section 2.1 in https://arxiv.org/abs/1801.09797
      quant_mask = tl.one_hot(quant_mask, self._n_elements_in_block)
      quant_mask = fastmath.stop_gradient(quant_mask)
      quant_mask += mask - fastmath.stop_gradient(mask)  # straight-through
      # We will sometimes (quant_prob of the batches) use the soft-mask instead
      # of the quantized mask to improve training stability (see paper above).
      select = fastmath.random.uniform(self.rng, (), jnp.float32, 0.0, 1.0)
      quant_mask = jnp.where(select < self._quant_prob, quant_mask, mask)

      # In training, run full matmul to get benefits from the above tricks.
      mid = jnp.einsum('bd,dxy->bxy', x, w1) * quant_mask
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      if self._multiply_by_controller_output:
        # We multiply only for quantized decisions, since for non-quantized
        # decisions we've already multiplied the output.
        mask_mult = jnp.where(select < self._quant_prob,
                              mask, jnp.ones_like(mask))
        # Stop-gradient is here, because we already have a pass-through gradient
        # (for quantized decisions).
        mask_mult = fastmath.stop_gradient(mask_mult)
        relu = relu * mask_mult
      res = jnp.einsum('bxy,yxd->bd', relu, w2) + b2
    elif self._mode == 'predict':
      # This implementation mimicks inference. It's not efficient for large
      # size of joint_batch, but at inference that will be 1 most of the time.
      # Shapes:
      # quant_mask is [joint_batch, self._d1]
      # w1 is [d_model, self._d1, self._d2]
      # we'll index w1 with advanced numpy indexing, first range over
      # self._d1 times the batch size, second range being quant_mask
      batch_size = quant_mask.shape[0]
      idx1 = jnp.array([jnp.arange(self._d1)] * batch_size)
      # flatten indices and select from w1
      idx1 = jnp.reshape(idx1, [-1])
      idx2 = jnp.reshape(quant_mask, [-1])
      w = w1[idx1, idx2, :]  # now we have per-element weights with batch dim
      w = jnp.reshape(w, [batch_size, self._d1, -1])
      mid = jnp.einsum('ai,aji->aj', x, w)
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      if self._multiply_by_controller_output:
        mask_mult = jnp.take_along_axis(mask, quant_mask[..., None], -1)[..., 0]
        relu = relu * mask_mult
      # w2 is [self._d1, self._d2, d_model]
      v = w2[idx1, idx2, :]
      v = jnp.reshape(v, [batch_size, self._d1, -1])
      res = jnp.einsum('ai,aij->aj', relu, v) + b2
    else:
      quant_mask = tl.one_hot(quant_mask, self._n_elements_in_block)
      mid = jnp.einsum('bd,dxy->bxy', x, w1) * quant_mask
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      if self._multiply_by_controller_output:
        relu = relu * mask
      res = jnp.einsum('bxy,yxd->bd', relu, w2) + b2

    return original_quant_mask, jnp.reshape(res, x_shape)
def forward ( self , x ) : if self . _mode == 'train' or self . _multiply_by_controller_output : quant_mask , mask , x = x else : quant_mask , x = x original_quant_mask = quant_mask w1 , w2 , b2 = self . weights if self . _mode == 'predict' : w1 = jnp . transpose ( w1 , ( 1 , 2 , 0 ) ) # dm, d1, d2 -> d1, d2, dm w2 = jnp . transpose ( w2 , ( 1 , 0 , 2 ) ) # d2, d1, dm -> d1, d2, dm x_shape = x . shape x = jnp . reshape ( x , [ - 1 , x_shape [ - 1 ] ] ) # Easier to operate on flattened x. if self . _mode == 'train' : # Tricks from Section 2.1 in https://arxiv.org/abs/1801.09797 quant_mask = tl . one_hot ( quant_mask , self . _n_elements_in_block ) quant_mask = fastmath . stop_gradient ( quant_mask ) quant_mask += mask - fastmath . stop_gradient ( mask ) # straight-through # We will sometimes (quant_prob of the batches) use the soft-mask instead # of the quantized mask to improve training stability (see paper above). select = fastmath . random . uniform ( self . rng , ( ) , jnp . float32 , 0.0 , 1.0 ) quant_mask = jnp . where ( select < self . _quant_prob , quant_mask , mask ) # In training, run full matmul to get benefits from the above tricks. mid = jnp . einsum ( 'bd,dxy->bxy' , x , w1 ) * quant_mask relu = jnp . where ( mid <= 0 , jnp . zeros_like ( mid ) , mid ) if self . _multiply_by_controller_output : # We multiply only for quantized decisions, since for non-quantized # decisions we've already multiplied the output. mask_mult = jnp . where ( select < self . _quant_prob , mask , jnp . ones_like ( mask ) ) # Stop-gradient is here, because we already have a pass-through gradient # (for quantized decisions). mask_mult = fastmath . stop_gradient ( mask_mult ) relu = relu * mask_mult res = jnp . einsum ( 'bxy,yxd->bd' , relu , w2 ) + b2 elif self . _mode == 'predict' : # This implementation mimicks inference. It's not efficient for large # size of joint_batch, but at inference that will be 1 most of the time. # Shapes: # quant_mask is [joint_batch, self._d1] # w1 is [d_model, self._d1, self._d2] # we'll index w1 with advanced numpy indexing, first range over # self._d1 times the batch size, second range being quant_mask batch_size = quant_mask . shape [ 0 ] idx1 = jnp . array ( [ jnp . arange ( self . _d1 ) ] * batch_size ) # flatten indices and select from w1 idx1 = jnp . reshape ( idx1 , [ - 1 ] ) idx2 = jnp . reshape ( quant_mask , [ - 1 ] ) w = w1 [ idx1 , idx2 , : ] # now we have per-element weights with batch dim w = jnp . reshape ( w , [ batch_size , self . _d1 , - 1 ] ) mid = jnp . einsum ( 'ai,aji->aj' , x , w ) relu = jnp . where ( mid <= 0 , jnp . zeros_like ( mid ) , mid ) if self . _multiply_by_controller_output : mask_mult = jnp . take_along_axis ( mask , quant_mask [ ... , None ] , - 1 ) [ ... , 0 ] relu = relu * mask_mult # w2 is [self._d1, self._d2, d_model] v = w2 [ idx1 , idx2 , : ] v = jnp . reshape ( v , [ batch_size , self . _d1 , - 1 ] ) res = jnp . einsum ( 'ai,aij->aj' , relu , v ) + b2 else : quant_mask = tl . one_hot ( quant_mask , self . _n_elements_in_block ) mid = jnp . einsum ( 'bd,dxy->bxy' , x , w1 ) * quant_mask relu = jnp . where ( mid <= 0 , jnp . zeros_like ( mid ) , mid ) if self . _multiply_by_controller_output : relu = relu * mask res = jnp . einsum ( 'bxy,yxd->bd' , relu , w2 ) + b2 return original_quant_mask , jnp . reshape ( res , x_shape )