class_docstr: Cross-replica Batch Normalization.
  
  At every step the full batch is used to calculate the batch statistics even
  within a distributed setting (note only with ``snt.(Tpu)Replicator``).
  
  See :class:`BaseBatchNorm` for details.
  
  Attributes:
    scale: If ``create_scale=True``, a trainable :tf:`Variable` holding the
      current scale after the module is connected for the first time.
    offset: If ``create_offset``, a trainable :tf:`Variable` holding the current
      offset after the module is connected for the first time.
code: def _moments(
      self, inputs: tf.Tensor, use_batch_stats: types.BoolLike
  ) -> Tuple[tf.Tensor, tf.Tensor]:
      replica_context = tf.distribute.get_replica_context()
      if replica_context is None:
          raise TypeError(
              "Cross replica batch norm cannot be called in cross-replica context."
          )
  
      if use_batch_stats:
          # Note: This uses var=E(x^2) - E(x)^2 instead of the more numerically
          # stable var=E((x-E(x))^2) as this means that with XLA the all_reduces can
          # be combined and a fusion removed giving significant speed-up.
          # If you see NaNs in your model please try the alternative formula and
          # file a bug with your use-case.
          mean = tf.reduce_mean(inputs, self._axis, keepdims=True)
          mean = replica_context.all_reduce("MEAN", mean)
          mean_of_squares = tf.reduce_mean(
              tf.square(inputs), self._axis, keepdims=True
          )
          mean_of_squares = replica_context.all_reduce(
              "MEAN", mean_of_squares
          )
          mean_squared = tf.square(mean)
          var = mean_of_squares - mean_squared
          return mean, var
  
      else:  # use moving statistics
          mean = self.moving_mean.value
          variance = self.moving_variance.value
          return mean, variance
  
file_path: sonnet/sonnet/src/distribute/distributed_batch_norm.py
imports: from typing import Optional, Tuple
  from sonnet.src import batch_norm
  from sonnet.src import initializers
  from sonnet.src import metrics
  from sonnet.src import once
  from sonnet.src import types
  import tensorflow
label_desc: Use image node_id 3 for calling the CrossReplicaBatchNorm obj's underlying member method code with example usage: obj._moments(inputs, use_batch_stats) and returns: mean,  var, mean,  variance