class_docstr: A dense (a.k.a. fully-connected, affine) layer.
  
  Dense layers are the prototypical example of a trainable layer, i.e., a layer
  with trainable weights. Each node in a dense layer computes a weighted sum of
  all node values from the preceding layer and adds to that sum a node-specific
  bias term. The full layer computation is expressed compactly in linear
  algebra as an affine map `y = Wx + b`, where `W` is a matrix and `y`, `x`,
  and `b` are vectors. The layer is trained, or "learns", by updating the
  values in `W` and `b`.
  
  Less commonly, a dense layer can omit the bias term and be a pure linear map:
  `y = Wx`.
class_docstr_tok: ["A","dense","(","a.k.a",".","fully-connected",",","affine",")","layer",".","Dense","layers","are","the","prototypical","example","of","a","trainable","layer",",","i.e.",",","a","layer","with","trainable","weights",".","Each","node","in","a","dense","layer","computes","a","weighted","sum","of","all","node","values","from","the","preceding","layer","and","adds","to","that","sum","a","node-specific","bias","term",".","The","full","layer","computation","is","expressed","compactly","in","linear","algebra","as","an","affine","map","`","y","=","Wx","+","b","`",",","where","`","W","`","is","a","matrix","and","`","y","`",",","`","x","`",",","and","`","b","`","are","vectors",".","The","layer","is","trained",",","or","``","learns","''",",","by","updating","the","values","in","`","W","`","and","`","b","`",".","Less","commonly",",","a","dense","layer","can","omit","the","bias","term","and","be","a","pure","linear","map",":","`","y","=","Wx","`","."]
docstr: 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.
docstr_tok: ["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","."]
code: 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
  
imports: from absl import logging
  import numpy
  import tensorflow
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
label_desc: Use image node_id 3 for calling the Dense obj's underlying member method code with example usage: obj.init_weights_and_state(input_signature) without return types