Reshape instead of view in TiledFusedLogitsLoss

[alanhuangyoo]

Follow-up to #8348, in the same file.

That PR moved TiledMLP.backward off view because the flatten of batch and sequence into one axis needs a copy when a caller hands in a non-contiguous activation. TiledFusedLogitsLoss.forward does the same flatten, under the same comment, and is still on view:

# flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1
x = x.view(-1, *x.shape[2:])
y = y.view(-1, *y.shape[2:])
if mask is not None:
    mask = mask.view(-1)

A transposed activation — one of the two layouts #8348's test covers — fails there:

RuntimeError: view size is not compatible with input tensor's size and stride
(at least one dimension spans across two contiguous subspaces)

What is and is not reachable

I ran the four input shapes through TiledFusedLogitsLoss.apply on master:

input	contiguous	master
x contiguous	yes	ok
x transposed	no	RuntimeError: view size is not compatible …
x channel slice	no	ok
y / mask strided	no	ok

Only the transposed case is reachable today. A hidden-dimension slice keeps the wider row stride, and that still admits merging batch into sequence, so it survives — same for the strided y and mask I tried. y and mask go through the same flatten and are changed with it rather than left on a spelling that happens to hold.

The unflatten at the end of forward stays a view: x_grad comes from zeros_like() of the already-flattened x, so it is contiguous by construction. Same reasoning for the x_grad.view(x_shape_orig) that #8348 left alone in TiledMLP.backward.

TiledLoss has no such flatten, so nothing to do there.

Test

TestTiledFusedLogitsLossInputLayout mirrors TestTiledMLPInputLayout and checks the loss against the same input made contiguous, so it pins the value and not just the absence of a throw.

On master:

FAILED tests/unit/ulysses_alst/test_tiled_compute.py::TestTiledFusedLogitsLossInputLayout::test_transposed_input_matches_a_contiguous_copy[2]
FAILED ...[4]
RuntimeError: view size is not compatible with input tensor's size and stride
2 failed

With this PR:

tests/unit/ulysses_alst/test_tiled_compute.py   14 passed
yapf --diff / flake8                           clean