code: def testWithTfFunctionTfArgs(self, autograph):
      if "TPU" in self.device_types:
          self.skipTest("Test not working on TPU")
          # TODO(tamaranorman) enable on TPU
  
      layer = batch_norm.BaseBatchNorm(
          moving_mean=TestMetric(),
          moving_variance=TestMetric(),
          create_scale=False,
          create_offset=False,
          data_format="NHWC",
      )
      layer = tf.function(layer, autograph=autograph)
  
      inputs = tf.ones([2, 5, 3, 3, 3])
      expected = tf.zeros_like(inputs)
  
      for is_training, use_batch_stats in itertools.product(
          (True, False), (True, False)
      ):
          # NOTE: The use of `tf.constant` means we require graph control flow
          outputs = layer(
              inputs,
              tf.constant(is_training),
              tf.constant(use_batch_stats),
          )
          self.assertAllEqual(outputs, expected)
  
imports: import itertools
  from absl.testing import parameterized
  import numpy
  from sonnet.src import batch_norm
  from sonnet.src import initializers
  from sonnet.src import test_utils
  import tensorflow
label_desc: Use image node_id 8 for calling the BaseBatchNormTest obj's underlying member method code with example usage: obj.testWithTfFunctionTfArgs(autograph) without return types