(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, except the final dimension
      is the layer's `filters` value, and the second to last dimension is
      shrinked if 'VALID' padding is used with kernel_size bigger than one.
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, except the final dimension
      is the layer's `filters` value, and the second to last dimension is
      shrinked if 'VALID' padding is used with kernel_size bigger than one.
    """
    if self._use_bias:
      if not isinstance(self.weights, (tuple, list)):
        raise ValueError(f'Weights should be a (w, b) tuple or list; '
                         f'instead got: {self.weights}')
      w, b = self.weights
    else:
      w = self.weights

    linear_results_before_shifting = jnp.einsum(
        '...lp,lkpd->...lkd', x, w)
    # TODO(jaszczur): this could be run after padding for better efficiency

    if self._kernel_size == 1:
      # With kernel size 1 we don't have to split or shift anything.
      linear_result = jnp.squeeze(linear_results_before_shifting, axis=-2)
    else:
      # We computed a result for every "pixel", but each direction from the
      # receptive field (there are 'self._kernel_size' such directions) must be
      # shifted by a different amount. The easiest way to do it is to split
      # the tensor to 'self._kernel_size' smaller tensors, shift each one
      # appropriately, and then sum them together.
      split_shifting_linear_results = jnp.split(
          linear_results_before_shifting, self._kernel_size, axis=-2)

      for i in range(self._kernel_size):
        # Each tensor has to be shifted a different amount.
        if self._padding == 'WRAP':
          # We can shift by padding and cutting. With 'wrap' padding we
          # essentially have a torus.
          padding = [(0, 0) for i in split_shifting_linear_results[i].shape]
          padding[-3] = ((self._kernel_size - 1) - i, i)
          split_shifting_linear_results[i] = jnp.pad(
              split_shifting_linear_results[i], padding, mode='wrap')
          split_shifting_linear_results[i] = split_shifting_linear_results[i][
              ..., (self._kernel_size-1)//2:-(self._kernel_size-1)//2, :, :]
        elif self._padding == 'SAME':
          # We can shift by padding and cutting.
          padding = [(0, 0) for i in split_shifting_linear_results[i].shape]
          padding[-3] = ((self._kernel_size - 1) - i, i)
          split_shifting_linear_results[i] = jnp.pad(
              split_shifting_linear_results[i], padding)
          split_shifting_linear_results[i] = split_shifting_linear_results[i][
              ..., (self._kernel_size-1)//2:-(self._kernel_size-1)//2, :, :]
          # TODO(jaszczur): improve efficiency by not padding things to cut
        elif self._padding == 'VALID':
          # We don't need to shift - just cut the leftmost and rightmost values.
          cut_left = (self._kernel_size - 1) - i
          cut_right = split_shifting_linear_results[i].shape[-3] - i
          split_shifting_linear_results[i] = split_shifting_linear_results[i][
              ..., cut_left:cut_right, :, :]
        else:
          raise ValueError(f'Invalid padding {self._padding}')
      # After shifting.
      shifted_linear_results = jnp.concatenate(split_shifting_linear_results,
                                               axis=-2)
      linear_result = jnp.sum(shifted_linear_results, axis=-2)

    if self._use_bias:
      return linear_result + b
    else:
      return linear_result
def forward ( self , x ) : if self . _use_bias : if not isinstance ( self . weights , ( tuple , list ) ) : raise ValueError ( f'Weights should be a (w, b) tuple or list; ' f'instead got: {self.weights}' ) w , b = self . weights else : w = self . weights linear_results_before_shifting = jnp . einsum ( '...lp,lkpd->...lkd' , x , w ) # TODO(jaszczur): this could be run after padding for better efficiency if self . _kernel_size == 1 : # With kernel size 1 we don't have to split or shift anything. linear_result = jnp . squeeze ( linear_results_before_shifting , axis = - 2 ) else : # We computed a result for every "pixel", but each direction from the # receptive field (there are 'self._kernel_size' such directions) must be # shifted by a different amount. The easiest way to do it is to split # the tensor to 'self._kernel_size' smaller tensors, shift each one # appropriately, and then sum them together. split_shifting_linear_results = jnp . split ( linear_results_before_shifting , self . _kernel_size , axis = - 2 ) for i in range ( self . _kernel_size ) : # Each tensor has to be shifted a different amount. if self . _padding == 'WRAP' : # We can shift by padding and cutting. With 'wrap' padding we # essentially have a torus. padding = [ ( 0 , 0 ) for i in split_shifting_linear_results [ i ] . shape ] padding [ - 3 ] = ( ( self . _kernel_size - 1 ) - i , i ) split_shifting_linear_results [ i ] = jnp . pad ( split_shifting_linear_results [ i ] , padding , mode = 'wrap' ) split_shifting_linear_results [ i ] = split_shifting_linear_results [ i ] [ ... , ( self . _kernel_size - 1 ) // 2 : - ( self . _kernel_size - 1 ) // 2 , : , : ] elif self . _padding == 'SAME' : # We can shift by padding and cutting. padding = [ ( 0 , 0 ) for i in split_shifting_linear_results [ i ] . shape ] padding [ - 3 ] = ( ( self . _kernel_size - 1 ) - i , i ) split_shifting_linear_results [ i ] = jnp . pad ( split_shifting_linear_results [ i ] , padding ) split_shifting_linear_results [ i ] = split_shifting_linear_results [ i ] [ ... , ( self . _kernel_size - 1 ) // 2 : - ( self . _kernel_size - 1 ) // 2 , : , : ] # TODO(jaszczur): improve efficiency by not padding things to cut elif self . _padding == 'VALID' : # We don't need to shift - just cut the leftmost and rightmost values. cut_left = ( self . _kernel_size - 1 ) - i cut_right = split_shifting_linear_results [ i ] . shape [ - 3 ] - i split_shifting_linear_results [ i ] = split_shifting_linear_results [ i ] [ ... , cut_left : cut_right , : , : ] else : raise ValueError ( f'Invalid padding {self._padding}' ) # After shifting. shifted_linear_results = jnp . concatenate ( split_shifting_linear_results , axis = - 2 ) linear_result = jnp . sum ( shifted_linear_results , axis = - 2 ) if self . _use_bias : return linear_result + b else : return linear_result