def forward(self, inputs):
    """Executes this layer as part of a forward pass through the model."""
    if self._n_in == 1:
      inputs = [inputs]
    new_inputs = []
    for old_input in inputs:
      batched_axes_shape = list(old_input.shape[:-self._n_last_axes_to_keep])
      batched_shape = [-1] + list(old_input.shape[-self._n_last_axes_to_keep:])
      new_inputs.append(jnp.reshape(old_input, batched_shape))
    new_inputs = tuple(new_inputs)
    if self._n_in == 1:
      new_inputs = new_inputs[0]
    res, layer_state = self.sublayer.pure_fn(
        new_inputs, self.weights[0], self.state[0], self.rng)
    self.state = (layer_state,)
    return jnp.reshape(res, batched_axes_shape + list(res.shape[1:]))
Executes this layer as part of a forward pass through the model.
trax/layers/combinators.py