def __call__(self, x, weights=None, state=None, rng=None):
    """Makes layers callable; for use in tests or interactive settings.

    This convenience method helps library users play with, test, or otherwise
    probe the behavior of layers outside of a full training environment. It
    presents the layer as callable function from inputs to outputs, with the
    option of manually specifying weights and non-parameter state per individual
    call. For convenience, weights and non-parameter state are cached per layer
    instance, starting from default values of `EMPTY_WEIGHTS` and `EMPTY_STATE`,
    and acquiring non-empty values either by initialization or from values
    explicitly provided via the weights and state keyword arguments, in which
    case the old weights will be preserved, and the state will be updated.

    Args:
      x: Zero or more input tensors, packaged as described in the `Layer` class
          docstring.
      weights: Weights or `None`; if `None`, use self's cached weights value.
      state: State or `None`; if `None`, use self's cached state value.
      rng: Single-use random number generator (JAX PRNG key), or `None`;
          if `None`, use a default computed from an integer 0 seed.

    Returns:
      Zero or more output tensors, packaged as described in the `Layer` class
      docstring.
    """
    weights = self.weights if weights is None else weights
    rng = self.rng if rng is None else rng
    if state is not None:
      self.state = state  # Needed if the model wasn't fully initialized.
    state = self.state
    outputs, new_state = self.pure_fn(x, weights, state, rng)
    self.state = new_state
    return outputs
Makes layers callable; for use in tests or interactive settings.

    This convenience method helps library users play with, test, or otherwise
    probe the behavior of layers outside of a full training environment. It
    presents the layer as callable function from inputs to outputs, with the
    option of manually specifying weights and non-parameter state per individual
    call. For convenience, weights and non-parameter state are cached per layer
    instance, starting from default values of `EMPTY_WEIGHTS` and `EMPTY_STATE`,
    and acquiring non-empty values either by initialization or from values
    explicitly provided via the weights and state keyword arguments, in which
    case the old weights will be preserved, and the state will be updated.

    Args:
      x: Zero or more input tensors, packaged as described in the `Layer` class
          docstring.
      weights: Weights or `None`; if `None`, use self's cached weights value.
      state: State or `None`; if `None`, use self's cached state value.
      rng: Single-use random number generator (JAX PRNG key), or `None`;
          if `None`, use a default computed from an integer 0 seed.

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