docs: def _build_statistics(self, input_batch, use_batch_stats, stat_dtype):
      """Builds the statistics part of the graph when using moving variance.
  
      Args:
        input_batch: Input batch Tensor.
        use_batch_stats: Boolean to indicate if batch statistics should be
          calculated, otherwise moving averages are returned.
        stat_dtype: TensorFlow datatype to use for the moving mean and variance.
  
      Returns:
        Tuple of (mean, variance), each of the same datatype as `input_batch`.
      """
      # Set up our moving statistics. When connecting in parallel, this is shared.
      if self.MOVING_MEAN not in self._initializers:
        self._initializers[self.MOVING_MEAN] = create_mean_initializer()
      self._moving_mean = tf.get_variable(
          "moving_mean",
          dtype=stat_dtype,
          shape=(self._num_channels,),
          collections=[
              tf.GraphKeys.MOVING_AVERAGE_VARIABLES,
              tf.GraphKeys.GLOBAL_VARIABLES,
          ],
          initializer=self._initializers[self.MOVING_MEAN],
          trainable=False)
  
      if self.MOVING_VARIANCE not in self._initializers:
        self._initializers[self.MOVING_VARIANCE] = create_variance_initializer()
      self._moving_variance = tf.get_variable(
          "moving_variance",
          dtype=stat_dtype,
          shape=(self._num_channels,),
          collections=[
              tf.GraphKeys.MOVING_AVERAGE_VARIABLES,
              tf.GraphKeys.GLOBAL_VARIABLES,
          ],
          initializer=self._initializers[self.MOVING_VARIANCE],
          trainable=False)
  
      def build_batch_stats():
        """Builds the batch statistics calculation ops."""
        mean, variance = tf.nn.moments(input_batch, self._axis,
                                       keep_dims=True, name="normalize_moments")
  
        return mean, variance
  
      def build_moving_stats():
        """Retrieves the moving statistics."""
        # If necessary, cast the moving statistics to match the input type.
        # This is required by tf.nn.batch_normalization.
        input_dtype = input_batch.dtype.base_dtype
        if stat_dtype == input_dtype:
          return (
              tf.identity(self._moving_mean),
              tf.identity(self._moving_variance),
          )
        else:
          return (
              tf.cast(self._moving_mean, input_dtype),
              tf.cast(self._moving_variance, input_dtype),
          )
  
      mean, variance = utils.smart_cond(
          use_batch_stats,
          build_batch_stats,
          build_moving_stats,
      )
  
      return mean, variance