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.
docstr: Constructs a ``CrossReplicaBatchNorm`` module.
  
  Args:
    create_scale: whether to create a trainable scale per channel applied
      after the normalization.
    create_offset: whether to create a trainable offset per channel applied
      after normalization and scaling.
    moving_mean: An object which keeps track of the moving average of the mean
      which can be used to normalize at test time. This object must have an
      update method which takes a value and updates the internal state and a
      value property which returns the current mean.
    moving_variance: An object which keeps track of the moving average of the
      variance which can be used to normalize at test time. This object must
      have an update method which takes a value and updates the internal state
      and a value property which returns the current variance.
    eps: Small epsilon to avoid division by zero variance. Defaults to
      ``1e-5``.
    scale_init: Optional initializer for the scale variable. Can only be set
      if ``create_scale=True``. By default scale is initialized to ``1``.
    offset_init: Optional initializer for the offset variable. Can only be set
      if ``create_offset=True``. By default offset is initialized to ``0``.
    data_format: The data format of the input. Can be either
      ``channels_first``, ``channels_last``, ``N...C`` or ``NC...``. By
      default it is ``channels_last``.
    name: Name of the module.
code: def __init__(
      self,
      create_scale: bool,
      create_offset: bool,
      moving_mean: metrics.Metric,
      moving_variance: metrics.Metric,
      eps: types.FloatLike = 1e-5,
      scale_init: Optional[initializers.Initializer] = None,
      offset_init: Optional[initializers.Initializer] = None,
      data_format: str = "channels_last",
      name: Optional[str] = None,
  ):
      """Constructs a ``CrossReplicaBatchNorm`` module.
  
      Args:
        create_scale: whether to create a trainable scale per channel applied
          after the normalization.
        create_offset: whether to create a trainable offset per channel applied
          after normalization and scaling.
        moving_mean: An object which keeps track of the moving average of the mean
          which can be used to normalize at test time. This object must have an
          update method which takes a value and updates the internal state and a
          value property which returns the current mean.
        moving_variance: An object which keeps track of the moving average of the
          variance which can be used to normalize at test time. This object must
          have an update method which takes a value and updates the internal state
          and a value property which returns the current variance.
        eps: Small epsilon to avoid division by zero variance. Defaults to
          ``1e-5``.
        scale_init: Optional initializer for the scale variable. Can only be set
          if ``create_scale=True``. By default scale is initialized to ``1``.
        offset_init: Optional initializer for the offset variable. Can only be set
          if ``create_offset=True``. By default offset is initialized to ``0``.
        data_format: The data format of the input. Can be either
          ``channels_first``, ``channels_last``, ``N...C`` or ``NC...``. By
          default it is ``channels_last``.
        name: Name of the module.
      """
      super().__init__(
          create_scale=create_scale,
          create_offset=create_offset,
          moving_mean=moving_mean,
          moving_variance=moving_variance,
          eps=eps,
          scale_init=scale_init,
          offset_init=offset_init,
          data_format=data_format,
          name=name,
      )
  
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 1 to create a new CrossReplicaBatchNorm object from inherited base classes: batch_norm with example: obj = CrossReplicaBatchNorm(create_scale, create_offset, moving_mean, moving_variance, eps, scale_init, offset_init, data_format, name)