Error in cross_entropy_logits_sparse function about objax HOT 3 CLOSED

def cross_entropy_logits_sparse(logits: JaxArray, labels: Union[JaxArray, int]) -> JaxArray:
    """Computes the softmax cross-entropy loss.
    Args:
        logits: (batch, ..., #class) tensor of logits.
        labels: (batch, ...) integer tensor of label indexes in {0, ...,#nclass-1} or just a single integer.
    Returns:
        (batch, ....) tensor of the cross-entropies for each entry.
    """
    if isinstance(labels, int):
        labeled_logits = logits[..., labels]
    else:
        labeled_logits = jn.take_along_axis(logits, labels[..., None], -1).squeeze(-1)
        
    return logsumexp(logits, axis=-1) - labeled_logits