from tensorflow.python.ops import state_ops

    else:
      batch_count = math_ops.reduce_sum(
          _broadcast_weights(weights, labels))  # n_B in eqn
      weighted_predictions = predictions * weights
      weighted_labels = labels * weights

    update_count = state_ops.assign_add(count, batch_count)  # n_AB in eqn
    prev_count = update_count - batch_count  # n_A in update equation

    # We update the means by Delta=Error*BatchCount/(BatchCount+PrevCount)
    # batch_mean_prediction is E[x_B] in the update equation
    batch_mean_prediction = _safe_div(
