def forward(self, inputs):
    """Overrides `Layer.forward`.

    Args:
      inputs: Zero or more input tensors, packaged as described in the `Layer`
          class docstring.

    Returns:
      Zero or more output tensors, packaged as described in the `Layer` class
      docstring.
    """
    _validate_forward_input(inputs, self.n_in)
    raw_output = self._forward_fn(inputs)
    output = () if _is_empty(raw_output) else raw_output
    return output
Overrides `Layer.forward`.

    Args:
      inputs: Zero or more input tensors, packaged as described in the `Layer`
          class docstring.

    Returns:
      Zero or more output tensors, packaged as described in the `Layer` class
      docstring.
    
trax/layers/base.py