code: def _save_batch_norm_model(export_dir, save_from_keras=False):
      """Writes a Hub-style SavedModel with a batch norm layer."""
      inp = tf_keras_v2.layers.Input(shape=(1,), dtype=tf.float32)
      bn = tf_keras_v2.layers.BatchNormalization(momentum=0.8)
      outp = bn(inp)
      model = tf_keras_v2.Model(inp, outp)
  
      if save_from_keras:
          tf.saved_model.save(model, export_dir)
          return
  
      @tf.function
      def call_fn(inputs, training=False):
          return model(inputs, training=training)
  
      for training in (True, False):
          call_fn.get_concrete_function(
              tf.TensorSpec((None, 1), tf.float32), training=training
          )
  
      obj = tf.train.Checkpoint()
      obj.__call__ = call_fn
      # Test assertions pick up variables by their position here.
      obj.trainable_variables = [bn.beta, bn.gamma]
      assert _tensors_names_set(
          obj.trainable_variables
      ) == _tensors_names_set(model.trainable_variables)
      obj.variables = [
          bn.beta,
          bn.gamma,
          bn.moving_mean,
          bn.moving_variance,
      ]
      assert _tensors_names_set(obj.variables) == _tensors_names_set(
          model.trainable_variables + model.non_trainable_variables
      )
      obj.regularization_losses = []
      assert not model.losses
      tf.saved_model.save(obj, export_dir)
  
imports: import json
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_hub
label_desc: Use image node_id 6 for calling a global function with example usage: _save_batch_norm_model(export_dir, save_from_keras) and returns: model