def forward(self, inputs):
    """Executes this layer as part of a forward pass through the model.

    Args:
      inputs: Tensors required by the sublayer.

    Returns:
      Tensors resulting from running the sublayer the first time.
    """
    state, weights = self.state, self.weights[0]
    if state[0] is ():  # pylint: disable=literal-comparison
      res, layer_state = self.sublayer.pure_fn(
          inputs, weights, state[1], self.rng)
      self.state = (res, layer_state)
      return res
    else:
      return state[0]
Executes this layer as part of a forward pass through the model.

    Args:
      inputs: Tensors required by the sublayer.

    Returns:
      Tensors resulting from running the sublayer the first time.
    
trax/layers/combinators.py