code: def custom_getter(self, activation_dtype=tf.bfloat16):
      
      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
docs: 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.