code: def get_variable_dtype(
      master_dtype=tf.bfloat16,
      slice_dtype=tf.float32,
      activation_dtype=tf.float32):
    
    return mtf.VariableDType(
        master_dtype=tf.as_dtype(master_dtype),
        slice_dtype=tf.as_dtype(slice_dtype),
        activation_dtype=tf.as_dtype(activation_dtype))
docs: Datatypes to use for the run.
  
    Args:
      master_dtype: string, datatype for checkpoints
        keep this the same between training and eval/inference
      slice_dtype: string, datatype for variables in memory
        must be tf.float32 for training
      activation_dtype: string, datatype for activations
        less memory usage if tf.bfloat16 but possible numerical issues
    Returns:
      a mtf.VariableDtype