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 with previously specified shape and dtype.
    """
    self.state = {}
    batch_size = x[0].shape[0]
    num_images = min(self._num_summaries, batch_size)
    for s in range(num_images):
      images = []
      for i in range(self._n_in):
        images.append(
            self._recover_fn(x[i][s]) if self._recover_fn else x[i][s])
      self.state[self._name + str(s)] = jnp.concatenate(images, axis=0)
    return x[:self._n_in]
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 with previously specified shape and dtype.
    
trax/layers/core.py