def weights_and_state_signature(self, input_signature, unsafe=False):
    """Return a pair containing the signatures of weights and state."""
    rng, state, weights = self.rng, self.state, self.weights
    abstract_init = fastmath.abstract_eval(self.init)
    sig = abstract_init(input_signature)
    self.rng = rng
    if not unsafe:
      self.state, self.weights = state, weights
    return sig
Return a pair containing the signatures of weights and state.
trax/layers/base.py