def __init__(self, initializer, shape=tuple(), use_bfloat16=False):
    """Returns a learnable tensor of shape `shape`.

    Args:
      initializer: Function taking shape and rng as arguments.
      shape: Shape of the learnable weights.
      use_bfloat16: If `True`, use bfloat16 weights instead of the default
        float32; this can save memory but may (rarely) lead to numerical issues.
    """
    super().__init__(name=f'Weights_{shape}', n_in=0, n_out=1)
    self._shape = shape
    self._initializer = initializer
    self._use_bfloat16 = use_bfloat16
Returns a learnable tensor of shape `shape`.

    Args:
      initializer: Function taking shape and rng as arguments.
      shape: Shape of the learnable weights.
      use_bfloat16: If `True`, use bfloat16 weights instead of the default
        float32; this can save memory but may (rarely) lead to numerical issues.
    
trax/layers/core.py