Introduction

tf.argmax() and tf.equal() often appear together when you convert model outputs into class labels and then measure whether those labels match expected values. They solve different problems: one chooses an index, and the other compares tensors element by element.

Using tf.argmax() to turn scores into labels

tf.argmax() returns the position of the largest value along a given axis. In classification code, that usually means turning a vector of scores into a class id.

1import tensorflow as tf
2
3logits = tf.constant([
4    [0.1, 2.4, 0.3],
5    [1.2, 0.7, 0.9],
6    [0.5, 0.6, 3.1],
7], dtype=tf.float32)
8
9predicted_classes = tf.argmax(logits, axis=1, output_type=tf.int32)
10print(predicted_classes.numpy())  # [1 0 2]

The axis argument matters. With a shape of (batch_size, num_classes), axis=1 selects the best class per row. If you accidentally use axis=0, TensorFlow finds the largest item in each column across the whole batch, which is rarely what you want for predictions.

wrong = tf.argmax(logits, axis=0)
print(wrong.numpy())  # column-wise result, not per-example classes

Another detail is output type. Labels are often stored as int32, while tf.argmax() defaults to int64 in many setups. If you plan to compare predictions with label tensors, matching the dtype early avoids unnecessary casting later.

Using tf.equal() to compare predictions and labels

tf.equal() performs an element-wise equality check and returns a boolean tensor. It does not calculate accuracy by itself, but it provides the raw comparison you need.

1import tensorflow as tf
2
3predicted_classes = tf.constant([1, 0, 2], dtype=tf.int32)
4true_classes = tf.constant([1, 2, 2], dtype=tf.int32)
5
6matches = tf.equal(predicted_classes, true_classes)
7print(matches.numpy())  # [ True False  True]

Once you have the boolean tensor, you can reduce it into an accuracy value.

accuracy = tf.reduce_mean(tf.cast(matches, tf.float32))
print(float(accuracy.numpy()))  # 0.6666666865348816

This pattern is common in quick experiments, custom training loops, and debugging sessions. It is especially useful when you want to inspect exactly which examples failed rather than relying only on a prebuilt metric object.

Putting both operations together in a model workflow

A typical sequence is:

1. Run the model and get logits or probabilities.
2. Use tf.argmax() on the class dimension.
3. Compare the result with labels using tf.equal().
4. Cast and reduce to compute a metric.

1import tensorflow as tf
2
3y_true = tf.constant([2, 1, 0], dtype=tf.int32)
4y_pred = tf.constant([
5    [0.2, 0.1, 0.7],
6    [0.3, 0.5, 0.2],
7    [0.9, 0.05, 0.05],
8], dtype=tf.float32)
9
10predicted = tf.argmax(y_pred, axis=1, output_type=tf.int32)
11correct = tf.equal(predicted, y_true)
12accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))
13
14print("predicted:", predicted.numpy())
15print("correct:", correct.numpy())
16print("accuracy:", float(accuracy.numpy()))

If your labels are one-hot encoded instead of integer class ids, convert them first:

1one_hot_labels = tf.constant([
2    [0, 0, 1],
3    [0, 1, 0],
4    [1, 0, 0],
5], dtype=tf.float32)
6
7label_ids = tf.argmax(one_hot_labels, axis=1, output_type=tf.int32)

That keeps the comparison shape aligned with the prediction tensor.

Common Pitfalls

The most common mistake is using the wrong axis in tf.argmax(). For batch-first classification output, axis=1 is usually correct, while axis=0 changes the meaning entirely.

Another frequent issue is comparing tensors with different dtypes. tf.equal() expects compatible tensors, so a prediction tensor of int64 and a label tensor of int32 may force extra casting or fail in stricter code paths. Set output_type explicitly or cast one side before comparing.

Shape mismatches also cause confusion. Comparing class ids shaped like (batch_size,) with one-hot labels shaped like (batch_size, num_classes) will not give the result you expect. Convert both tensors into the same representation first.

Finally, remember that tf.equal() is not a ranking or tolerance-based check. For floating-point probabilities, exact equality is usually the wrong tool. Compare class ids or use a numeric tolerance method when you are validating floats.

Summary

• Use tf.argmax() to convert a score vector into a class index.
• Choose the axis carefully; for (batch_size, num_classes), axis=1 is the usual choice.
• Use tf.equal() for element-wise comparison after predictions and labels share the same shape and dtype.
• Cast boolean results to float32 and reduce them to compute accuracy.
• Convert one-hot labels to class ids before comparing them with argmax-based predictions.