class_docstr: Cross replica batch normalization.
class_docstr_tok: ["Cross","replica","batch","normalization","."]
docstr: Compute the mean and variance: it overrides the original _moments.
docstr_tok: ["Compute","the","mean","and","variance",":","it","overrides","the","original","_moments","."]
code: def _moments(self, inputs, reduction_axes, keep_dims):
      """Compute the mean and variance: it overrides the original _moments."""
      shard_mean, shard_variance = super(
          TpuBatchNormalization, self
      )._moments(inputs, reduction_axes, keep_dims=keep_dims)
  
      num_shards = tpu_function.get_tpu_context().number_of_shards or 1
      if (
          num_shards <= 8
      ):  # Skip cross_replica for 2x2 or smaller slices.
          num_shards_per_group = 1
      else:
          num_shards_per_group = max(8, num_shards // 8)
      logging.info(
          "TpuBatchNormalization with num_shards_per_group %s",
          num_shards_per_group,
      )
      if num_shards_per_group > 1:
          # Compute variance using: Var[X]= E[X^2] - E[X]^2.
          shard_square_of_mean = tf.math.square(shard_mean)
          shard_mean_of_square = shard_variance + shard_square_of_mean
          group_mean = self._cross_replica_average(
              shard_mean, num_shards_per_group
          )
          group_mean_of_square = self._cross_replica_average(
              shard_mean_of_square, num_shards_per_group
          )
          group_variance = group_mean_of_square - tf.math.square(
              group_mean
          )
          return (group_mean, group_variance)
      else:
          return (shard_mean, shard_variance)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import json
  import os
  import sys
  from absl import logging
  import numpy
  import tensorflow.compat.v1
  from tensorflow.python.tpu import tpu_function
label_desc: Use image node_id 3 for calling the TpuBatchNormalization obj's underlying member method code with example usage: obj._moments(inputs, reduction_axes, keep_dims) and returns: group_mean,  group_variance, shard_mean,  shard_variance