class_docstr: Learnable weights as a layer.
  
  It takes no input and returns a single tensor: weights.
class_docstr_tok: ["Learnable","weights","as","a","layer",".","It","takes","no","input","and","returns","a","single","tensor",":","weights","."]
docstr: 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.
docstr_tok: ["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","."]
code: 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
  
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 1 to create a new Weights object from inherited base classes: base with example: obj = Weights(initializer, shape, use_bfloat16)