class_docstr: Cross replica batch normalization.
class_docstr_tok: ["Cross","replica","batch","normalization","."]
docstr: Calculates the average value of input tensor across TPU replicas.
docstr_tok: ["Calculates","the","average","value","of","input","tensor","across","TPU","replicas","."]
code: def _cross_replica_average(self, t, num_shards_per_group):
      """Calculates the average value of input tensor across TPU replicas."""
      num_shards = tpu_function.get_tpu_context().number_of_shards
      group_assignment = None
      if num_shards_per_group > 1:
          if num_shards % num_shards_per_group != 0:
              raise ValueError(
                  "num_shards: %d mod shards_per_group: %d, should be 0"
                  % (num_shards, num_shards_per_group)
              )
          num_groups = num_shards // num_shards_per_group
          group_assignment = [
              [
                  x
                  for x in range(num_shards)
                  if x // num_shards_per_group == y
              ]
              for y in range(num_groups)
          ]
      return tf.tpu.cross_replica_sum(t, group_assignment) / tf.cast(
          num_shards_per_group, t.dtype
      )
  
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 2 for calling the TpuBatchNormalization obj's underlying member method code with example usage: obj._cross_replica_average(t, num_shards_per_group) and returns: unknown