def init_weights_and_state(self, input_signature):
    """Initializes this layer's single weight to zero."""
    del input_signature
    self.weights = jnp.zeros((), dtype=jnp.float32)
Initializes this layer's single weight to zero.
trax/layers/activation_fns.py