def __init__(self, forward_fn, n_in=1, n_out=1, name='PureLayer'):
    """Creates an unconnected `PureLayer` instance.

    Args:
      forward_fn: Pure function from input tensors to output tensors, where
          inputs and outputs are packaged as specified for `forward`.
      n_in: Number of inputs expected by this layer.
      n_out: Number of outputs promised by this layer.
      name: Class-like name for this layer; for use only in debugging.
    """
    super().__init__(n_in, n_out, name)
    self._forward_fn = forward_fn
Creates an unconnected `PureLayer` instance.

    Args:
      forward_fn: Pure function from input tensors to output tensors, where
          inputs and outputs are packaged as specified for `forward`.
      n_in: Number of inputs expected by this layer.
      n_out: Number of outputs promised by this layer.
      name: Class-like name for this layer; for use only in debugging.
    
trax/layers/base.py