import tensorflow as tf

    # `is_training` is unknown.
    is_training_const = utils.constant_value(is_training)
    if is_training_const is None or is_training_const:
      update_mean_op, update_second_moment_op = utils.smart_cond(
          is_training,
          build_update_ops,
          build_no_ops,
      )

      # Every new connection creates a new op which adds its contribution
      # to the running average when ran.
      tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_mean_op)
      tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_second_moment_op)

  def _build(self, input_batch, is_training=True, test_local_stats=True):
    """Connects the BatchNorm module into the graph.

    Args:
      input_batch: A Tensor of arbitrary dimension. By default, the final
        dimension is not reduced over when computing the minibatch statistics.
      is_training: A boolean to indicate if the module should be connected in
        training mode, meaning the moving averages are updated. By default
        `True`. Can be a Tensor.
      test_local_stats: A boolean to indicate if local batch statistics should
