def init_weights_and_state(self, input_signature):
    """Randomly initializes this layer's weights.

    Weights are a `(w, b)` tuple for layers created with `use_bias=True` (the
    default case), or a `w` tensor for layers created with `use_bias=False`.

    Args:
      input_signature: `ShapeDtype` instance characterizing the input this layer
          should compute on.
    """
    shape_w = (input_signature.shape[-1], self._n_units)
    shape_b = (self._n_units,)
    rng_w, rng_b = fastmath.random.split(self.rng, 2)
    w = self._kernel_initializer(shape_w, rng_w)
    if self._use_bfloat16:
      w = w.astype(jnp.bfloat16)

    if self._use_bias:
      b = self._bias_initializer(shape_b, rng_b)
      if self._use_bfloat16:
        b = b.astype(jnp.bfloat16)
      self.weights = (w, b)
    else:
      self.weights = w
Randomly initializes this layer's weights.

    Weights are a `(w, b)` tuple for layers created with `use_bias=True` (the
    default case), or a `w` tensor for layers created with `use_bias=False`.

    Args:
      input_signature: `ShapeDtype` instance characterizing the input this layer
          should compute on.
    
trax/layers/core.py