Fix InputNormalization padding masks for nondefault length dimensions - #3109

What does this PR do?

InputNormalization uses length_dim to compute statistics but always builds its padding mask along dimension 1. For channel-first features, length_dim=2 therefore gives wrong sentence-normalized values and raises AssertionError with batch/global normalization.

This reproducer compares the same padded features in time-first and channel-first layouts:

import torch
from speechbrain.processing.features import InputNormalization

x = torch.tensor([
    [[1.0, 10.0], [3.0, 14.0], [5.0, 18.0], [7.0, 22.0]],
    [[2.0, 20.0], [4.0, 24.0], [0.0, 0.0], [0.0, 0.0]],
])
lengths = torch.tensor([1.0, 0.5])
for norm_type in ("sentence", "batch", "global"):
    expected = InputNormalization(
        norm_type=norm_type, avoid_padding_norm=True
    )(x, lengths)
    norm = InputNormalization(
        norm_type=norm_type, length_dim=2, avoid_padding_norm=True
    )
    try:
        actual = norm(x.transpose(1, 2), lengths).transpose(1, 2)
        print(
            f"{norm_type}: matches_time_first={torch.allclose(actual, expected)}, "
            f"short_utterance_first_feature={actual[1, :2, 0].tolist()}"
        )
    except AssertionError:
        print(f"{norm_type}: AssertionError")

Observed on CPU with torch/torchaudio 2.11.0:

I pass self.length_dim to the existing mask helper and clarify the argument's docstring. The regression test moves the time axis to dimensions 2 and 3 for all three normalization modes, with and without padding-output normalization. All 12 cases fail before the fix and pass after it.

This affects callers using the documented nondefault time axis, such as a channel-first spectrogram pipeline. I did not find a bundled recipe using a nondefault length_dim. The default layout is unchanged. No new dependencies or breaking API changes.

Checks run:

• pytest tests/unittests/test_input_norm.py tests/unittests/test_features.py tests/unittests/test_length_mask.py --doctest-modules speechbrain/processing/features.py -q: 564 passed, 1 GPU-only test skipped.
• Ruff 0.12.4 lint and format checks passed for both changed files.
• pre-commit run -a: all hooks passed.

PR review

I used AI assistance (Claude and GPT-6 Astra) to find, fix and test this bug.