class_docstr: Helper class for encoding weights as bfloat16.
  
  For now, the parameters are always stored (encoded) as bfloat16 and decoded
  to bfloat32.  Confusingly, the custom getter then converts the bfloat32 back
  to a bfloat16 to use as an activation, assuming that we use bfloat16 for
  activations.
  
  TODO(noam): Add options for activation dtype=float32, and for different
  storage dtypes.
docstr: A custom getter that uses the encoding for bfloat16 and float32 vars.
  
  When a bfloat16 or float32 variable is requsted, an encoded float16
  varaible is created, which is then decoded and cast to a bfloat16
  activation.
  
  Args:
    activation_dtype: a dtype to which to convert the decoded value.
  
  Returns:
    a function.
docstr_tok: ["A","custom","getter","that","uses","the","encoding","for","bfloat16","and","float32","vars",".","When","a","bfloat16","or","float32","variable","is","requsted",",","an","encoded","float16","varaible","is","created",",","which","is","then","decoded","and","cast","to","a","bfloat16","activation",".","Args",":","activation_dtype",":","a","dtype","to","which","to","convert","the","decoded","value",".","Returns",":","a","function","."]
code: def custom_getter(self, activation_dtype=tf.bfloat16):
      """A custom getter that uses the encoding for bfloat16 and float32 vars.
  
      When a bfloat16 or float32 variable is requsted, an encoded float16
      varaible is created, which is then decoded and cast to a bfloat16
      activation.
  
      Args:
        activation_dtype: a dtype to which to convert the decoded value.
  
      Returns:
        a function.
      """
  
      def getter_fn(getter, *args, **kwargs):
          requested_dtype = kwargs["dtype"]
          if requested_dtype in (tf.bfloat16, tf.float32):
              kwargs["dtype"] = tf.bfloat16
              kwargs["initializer"] = _EncodingInitializer(
                  kwargs["initializer"], self
              )
              ret = self._decode_with_identity_gradient(
                  getter(*args, **kwargs)
              )
              return tf.cast(ret, activation_dtype)
          return getter(*args, **kwargs)
  
      return getter_fn
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow.compat.v1
  from tensorflow.python.framework import function
label_desc: Use image node_id 4 for calling the ParameterEncoding obj's underlying member method code with example usage: obj.custom_getter(activation_dtype) and returns: getter_fn, getter, tf