def pure_fn(self, x, weights, state, rng, use_cache=False):
    """Applies this layer as a pure function with no optional args.

    This method exposes the layer's computation as a pure function. This is
    especially useful for JIT compilation. Do not override, use `forward`
    instead.

    Args:
      x: Zero or more input tensors, packaged as described in the `Layer` class
          docstring.
      weights: A tuple or list of trainable weights, with one element for this
          layer if this layer has no sublayers, or one for each sublayer if
          this layer has sublayers. If a layer (or sublayer) has no trainable
          weights, the corresponding weights element is an empty tuple.
      state: Layer-specific non-parameter state that can update between batches.
      rng: Single-use random number generator (JAX PRNG key).
      use_cache: if `True`, cache weights and state in the layer object; used
        to implement layer sharing in combinators.

    Returns:
      A tuple of `(tensors, state)`. The tensors match the number (`n_out`)
      promised by this layer, and are packaged as described in the `Layer`
      class docstring.
    """
    try:
      old_weights, old_state, old_rng = self.weights, self.state, self.rng
      self._rng = rng
      # The isinstance check is only needed when == is overloaded, as in TF.
      if (isinstance(weights, dict) and isinstance(state, dict) and
          weights == GET_WEIGHTS_FROM_CACHE and state == GET_STATE_FROM_CACHE):
        was_cached = True
        weights = self.weights
        state = self.state
      else:
        # In this case, we're called for the first time: cache weights.
        was_cached = False
        self.weights, self.state = weights, state

      # If weights are sharded across multiple devices, unshard before forward.
      sharded_weights, weights_were_unsharded = weights, False
      if N_WEIGHTS_SHARDS > 1 and not self.sublayers:
        self.weights, weights_were_unsharded = unshard_in_pmap(
            weights, N_WEIGHTS_SHARDS)

      if not self.has_backward:
        outputs = self.forward(x)
        s = self.state
      else:
        outputs, s = self._do_custom_gradients(x)
        self.state = s
      self._rng = old_rng
      if weights_were_unsharded:  # only store a shard of weights if sharded
        self.weights = sharded_weights

      if not use_cache:
        self.weights, self.state = old_weights, old_state
      if was_cached:  # If the layer was shared, return a state marking this.
        s = GET_STATE_FROM_CACHE
      return outputs, s

    except Exception:
      # Skipping 3 lines as it's always the uninteresting internal call.
      name, trace = self._name, _short_traceback(skip=3)
      raise LayerError(name, 'pure_fn',
                       self._caller, signature(x), trace) from None
Applies this layer as a pure function with no optional args.

    This method exposes the layer's computation as a pure function. This is
    especially useful for JIT compilation. Do not override, use `forward`
    instead.

    Args:
      x: Zero or more input tensors, packaged as described in the `Layer` class
          docstring.
      weights: A tuple or list of trainable weights, with one element for this
          layer if this layer has no sublayers, or one for each sublayer if
          this layer has sublayers. If a layer (or sublayer) has no trainable
          weights, the corresponding weights element is an empty tuple.
      state: Layer-specific non-parameter state that can update between batches.
      rng: Single-use random number generator (JAX PRNG key).
      use_cache: if `True`, cache weights and state in the layer object; used
        to implement layer sharing in combinators.

    Returns:
      A tuple of `(tensors, state)`. The tensors match the number (`n_out`)
      promised by this layer, and are packaged as described in the `Layer`
      class docstring.
    
trax/layers/base.py