code: def test_quantize_bf16(self):
      # for custom model, quantized model still return fp32 output
      # test at the same time that the quantization does not
      # change the original's model dtype policy
      model = MyModel(100)
      model.compile(loss="mse", metrics=MeanSquaredError())
      ori_model_policies = []
      for layer in model.layers:
          ori_model_policies.append(layer._dtype_policy)
      x = np.random.random((100, 4))
      model(x)
  
      bf16_model = InferenceOptimizer.quantize(model, precision="bf16")
      for idx, layer in enumerate(model.layers):
          assert layer._dtype_policy == ori_model_policies[idx]
  
      from bigdl.nano.utils.common import _avx512_checker
  
      if _avx512_checker():
          output = bf16_model(x)
          assert output.dtype == tf.float32
  
      with tempfile.TemporaryDirectory() as tmp_dir_name:
          InferenceOptimizer.save(bf16_model, tmp_dir_name)
          load_model = InferenceOptimizer.load(tmp_dir_name, model)
  
      if _avx512_checker():
          output = load_model(x)
          assert output.dtype == tf.float32
  
      # test standard model, quantized model return bf16 output
      # test at the same time that the quantization does not
      # change the original's model dtype policy
      model = MobileNetV2(weights="imagenet")
      ori_model_config = model.get_config()
      x = np.random.rand(32, 224, 224, 3)
      model(x)
  
      bf16_model = InferenceOptimizer.quantize(model, precision="bf16")
      assert ori_model_config == model.get_config()
  
      from bigdl.nano.utils.common import _avx512_checker
  
      if _avx512_checker():
          output = bf16_model(x)
          assert output.dtype == tf.bfloat16
  
      with tempfile.TemporaryDirectory() as tmp_dir_name:
          InferenceOptimizer.save(bf16_model, tmp_dir_name)
          load_model = InferenceOptimizer.load(tmp_dir_name, model)
  
      if _avx512_checker():
          output = load_model(x)
          assert output.dtype == tf.bfloat16
  
      # test model with ModelWithConv2DTranspose layer, such model
      # cannot be saved after modifying dtype policy
      model = ModelWithConv2DTranspose()
      bf16_model = InferenceOptimizer.quantize(model, precision="bf16")
      inputs = tf.random.normal((1, 32, 32, 3))
      if _avx512_checker():
          output1 = model(inputs)
          output2 = bf16_model(inputs)
          assert output1.dtype == tf.float32
          assert output2.dtype == tf.bfloat16
          np.testing.assert_allclose(
              output1, tf.cast(output2, tf.float32), atol=1e-2
          )
  
imports: from unittest import TestCase
  import tempfile
  import operator
  import numpy
  import tensorflow
  import pytest
  from tensorflow.keras.metrics import MeanSquaredError, CategoricalAccuracy
  from tensorflow.keras import layers, Model
  from tensorflow.keras.applications import MobileNetV2
  from bigdl.nano.tf.keras import InferenceOptimizer
label_desc: Use image node_id 4 for calling the TestTraceAndQuantize obj's underlying member method code with example usage: obj.test_quantize_bf16() without return types