code: def enable_save_as_bf16(variables: List[tf_variables.Variable]):
      """Allows float32 DVariables to be checkpointed and restored as bfloat16.
  
      The method only affects the DVariable part inside the model and leaves
      non-DTensor Variables/Tensors untouched.
  
      Args:
        variables: A list of tf.Variable to be enabled with bfloat16 save/restore.
          Only has effect on DTensor Variables as they go through d_variables with
          DTensor Specific logis.
      """
      for v in variables:
          if isinstance(v, d_variable.DVariable):
              v.save_as_bf16 = True
  
file_path: tensorflow/tensorflow/dtensor/python/save_restore.py
imports: import collections
  from typing import Dict, List, Union
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import gen_dtensor_ops
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import mesh_util
  from tensorflow.python.eager import context
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor
  from tensorflow.python.ops import io_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 2 for calling a global function with example usage: enable_save_as_bf16(variables) without return types