Support tuple of axis in softmax_cross_entropy_with_integer_labels

[daskol]

Fix #1162 (see #1164 for concurrent fix).

[rdyro]

Hey, thanks! These PR looks great, I left some comments!

[rdyro]

JAX has an example of canonicalize_axis, we could vendor it from there (but it's a private API, so let's not import it from JAX directly)

[daskol]

Yeah, I know of canonicalize_axis and decided to get normalize_axis_* from origin for multiple reasons. Firstly, normilize_axis_* are proper implementation. Secondly, it basically consists of list/tuple transformation and jit-compilable by design. Thirdly, NumPy is already among dependencies and there is an explicit notice about deprecation in NumPy 2. But vendoring these routines is nice too.

[rdyro]

Yeah, I know of canonicalize_axis and decided to get normalize_axis_* from origin for multiple reasons. Firstly, normilize_axis_* are proper implementation. Secondly, it basically consists of list/tuple transformation and jit-compilable by design. Thirdly, NumPy is already among dependencies and there is an explicit notice about deprecation in NumPy 2. But vendoring these routines is nice too.

Makes sense, numpy's canonicalize are nice, but the import discrepancy between versions is a bit of a downside given potential API instability. Thanks for vendoring them!