softmax_cross_entropy_with_integer_labels

softmax_cross_entropy_with_integer_labels#

class braintools.metric.softmax_cross_entropy_with_integer_labels(logits, labels)#

Compute softmax cross entropy between logits and integer labels.

This is a more efficient version of softmax cross entropy when labels are provided as integer class indices rather than one-hot encoded vectors. Measures the probability error in discrete classification tasks where the classes are mutually exclusive.

Parameters:

  • logits (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Unnormalized log probabilities with shape [..., num_classes]. Must be floating point type.

  • labels (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Integer class indices specifying the correct class for each input. Values should be in the range [0, num_classes). Shape [...]. Must be integer type.

Returns:

Cross entropy between each prediction and the corresponding target distributions, with shape [...].

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

Examples

>>> import jax.numpy as jnp
>>> import braintools
>>> logits = jnp.array([[2.0, 1.0, 0.1]])
>>> labels = jnp.array([0])  # Class 0
>>> loss = braintools.metric.softmax_cross_entropy_with_integer_labels(logits, labels)
>>> print(loss)
[0.41702995]

Notes

Labels are used to gather along the last axis via jnp.take_along_axis. Out-of-range integer labels (outside [0, num_classes)) are not validated: JAX silently clamps out-of-bounds gather indices to the nearest valid index, so an invalid label yields a (silently) wrong loss rather than an error. Ensure labels lie within [0, num_classes).

References