def init(self, input_signature, rng=None, use_cache=False):
    """Initializes weights/state of this layer and its sublayers recursively.

    Initialization creates layer weights and state, for layers that use them.
    It derives the necessary array shapes and data types from the layer's input
    signature, which is itself just shape and data type information.

    For layers without weights or state, this method safely does nothing.

    This method is designed to create weights/state only once for each layer
    instance, even if the same layer instance occurs in multiple places in the
    network. This enables weight sharing to be implemented as layer sharing.

    Args:
      input_signature: `ShapeDtype` instance (if this layer takes one input)
          or list/tuple of `ShapeDtype` instances.
      rng: Single-use random number generator (JAX PRNG key), or `None`;
          if `None`, use a default computed from an integer 0 seed.
      use_cache: If `True`, and if this layer instance has already been
          initialized elsewhere in the network, then return special marker
          values -- tuple `(GET_WEIGHTS_FROM_CACHE, GET_STATE_FROM_CACHE)`.
          Else return this layer's newly initialized weights and state.

    Returns:
      A `(weights, state)` tuple.
    """
    try:
      if self._init_cached and use_cache:
        return (GET_WEIGHTS_FROM_CACHE, GET_STATE_FROM_CACHE)

      if rng is not None:
        self.rng = rng
      self.init_weights_and_state(input_signature)

      if use_cache:
        self._init_cached = True
      else:
        self._clear_init_cache()

      return (self.weights, self.state)

    except Exception:
      # Skipping 3 lines as it's always the uninteresting internal call.
      name, trace = self._name, _short_traceback(skip=3)
      raise LayerError(name, 'init', self._caller,
                       input_signature, trace) from None
Initializes weights/state of this layer and its sublayers recursively.

    Initialization creates layer weights and state, for layers that use them.
    It derives the necessary array shapes and data types from the layer's input
    signature, which is itself just shape and data type information.

    For layers without weights or state, this method safely does nothing.

    This method is designed to create weights/state only once for each layer
    instance, even if the same layer instance occurs in multiple places in the
    network. This enables weight sharing to be implemented as layer sharing.

    Args:
      input_signature: `ShapeDtype` instance (if this layer takes one input)
          or list/tuple of `ShapeDtype` instances.
      rng: Single-use random number generator (JAX PRNG key), or `None`;
          if `None`, use a default computed from an integer 0 seed.
      use_cache: If `True`, and if this layer instance has already been
          initialized elsewhere in the network, then return special marker
          values -- tuple `(GET_WEIGHTS_FROM_CACHE, GET_STATE_FROM_CACHE)`.
          Else return this layer's newly initialized weights and state.

    Returns:
      A `(weights, state)` tuple.
    
trax/layers/base.py