(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.
Applies this layer as a pure function with no optional args.
Applies this layer as a pure function with no optional args .
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
def pure_fn ( self , x , weights , state , rng , use_cache = False ) : 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