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
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.
    
trax/layers/core.py