Softmax (reuses Part 1, the first cross-part reuse)

The problem, from first principles

[01-classical-ml/02-classification/06-softmax-cce]'s softmax, all the way back in this curriculum's very first Part, was written for exactly one shape of input: a (batch_size, num_classes) matrix, softmax applied along axis=1 (each ROW normalized into a probability distribution over classes). [01-scaled-dot-product-attention]'s attention scores have a fundamentally different, and more VARIABLE, shape: (..., seq_len_q, seq_len_k), with anywhere from zero to several leading batch and head dimensions depending on the calling context (a single unbatched sequence, a batch of sequences, [04-multi-head-attention-split]'s per-head batched computation with an EXTRA head dimension on top of the batch dimension). Softmax always needs to normalize along the LAST axis of whatever that shape happens to be, not a fixed axis=1.

Rather than write an entirely new softmax implementation from scratch, this question demonstrates the same reuse-across-parts principle this whole curriculum has built toward, from [00-math-and-statistics]'s foundational math questions being reused throughout [01-classical-ml], to [02-deep-learning-core]'s autograd engine underlying everything in [03-dl-training], applied here for the FIRST time across an entire Part boundary: Part 1's softmax, entirely UNCHANGED, gets reused directly to build this more general version, via a simple reshaping trick rather than any actual reimplementation of the softmax math itself.

From theory to code

Implement softmax_last_axis(Z). Reshape Z (whatever its original rank) down to a 2D array, (-1, Z.shape[-1]), collapsing every leading dimension into one, so the LAST axis of the original array becomes exactly axis=1 of the reshaped 2D array, precisely the axis Part 1's softmax already operates on. Call Part 1's softmax (already provided, reused via load_solution) UNCHANGED on this reshaped 2D array, then reshape the result back to Z's original shape.

Constraints

• Must call Part 1's softmax function DIRECTLY (via the already-provided softmax_axis1 import), not reimplement the softmax formula independently.
• The reshape-and-reshape-back round trip must preserve the ORIGINAL shape of Z exactly, for any rank (2D, 3D, 4D, ...).
• Every "row" along the last axis (in the original, un-reshaped sense) must sum to exactly 1.

Hints