Output layer — predict the class

Background

A multi-class network's output layer produces one score (logit or probability) per class. To turn that into a prediction, pick the class with the highest score — the argmax. Softmax is monotonic, so the argmax of the logits equals the argmax of the probabilities; you don't even need to normalise to choose the winner.

Problem statement

Implement predict_class(scores) returning the predicted class index for each example.

Input

• scores — array-like of shape (batch, K): per-class scores (logits or probabilities), one row per example.

Output

Returns an np.ndarray of batch integer class indices (the column of the max in each row).

Examples

Example 1

Input:  scores = [[1, 3, 2], [5, 0, 1]]
Output: [1 0]

Explanation: row 0's max is at index 1; row 1's max is at index 0.

Example 2 — ties pick the first max

Input:  scores = [[1, 1, 0]]
Output: [0]

Explanation: indices 0 and 1 tie; argmax returns the first.

Constraints

• Take the argmax along the class axis (axis=1).
• Return an integer np.ndarray of length batch.

Notes

• Because softmax preserves order, classifying from logits gives the same labels as classifying from softmax probabilities — handy when you only need the decision, not the calibrated probability.