tests/collections/audio/test_audio_losses.py

# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

import numpy as np
import pytest
import torch

from nemo.collections.audio.losses.audio import (
    MAELoss,
    MSELoss,
    SDRLoss,
    calculate_mae_batch,
    calculate_mean,
    calculate_mse_batch,
    calculate_sdr_batch,
    convolution_invariant_target,
    scale_invariant_target,
)

try:
    import importlib

    importlib.import_module('torchaudio')

    HAVE_TORCHAUDIO = True
except ModuleNotFoundError:
    HAVE_TORCHAUDIO = False

from nemo.collections.audio.losses.maxine import CombinedLoss
from nemo.collections.audio.parts.utils.audio import (
    calculate_sdr_numpy,
    convolution_invariant_target_numpy,
    scale_invariant_target_numpy,
)


class TestAudioLosses:
    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('use_mask', [True, False])
    @pytest.mark.parametrize('use_input_length', [True, False])
    def test_calculate_mean(self, num_channels: int, use_mask: bool, use_input_length: bool):
        """Test mean calculation"""
        batch_size = 8
        num_samples = 50
        num_batches = 10
        random_seed = 42
        eps = 1e-10
        atol = 1e-6

        rng = torch.Generator()
        rng.manual_seed(random_seed)

        for n in range(num_batches):

            input_signal = torch.randn(size=(batch_size, num_channels, num_samples), generator=rng)
            # Random input length
            input_length = torch.randint(low=1, high=num_samples, size=(batch_size,), generator=rng)
            # Corresponding mask
            mask = torch.zeros(size=(batch_size, 1, num_samples))
            for i in range(batch_size):
                mask[i, :, : input_length[i]] = 1.0

            if use_mask and use_input_length:
                with pytest.raises(RuntimeError):
                    calculate_mean(input_signal, mask=mask, input_length=input_length, eps=eps)
                # Done with this test
                continue

            if use_mask:
                uut = calculate_mean(input_signal, mask=mask, eps=eps)
            elif use_input_length:
                uut = calculate_mean(input_signal, input_length=input_length, eps=eps)
            else:
                uut = calculate_mean(input_signal, eps=eps)

            # Calculate mean manually
            for b in range(batch_size):
                for c in range(num_channels):
                    golden = torch.mean(
                        input_signal[b, c, : input_length[b] if use_input_length or use_mask else num_samples]
                    )
                    assert torch.allclose(
                        uut[b, c], golden, atol=atol
                    ), f"Mean not matching for example {n}, channel {c}"

    @pytest.mark.unit
    def test_calculate_sdr_scale_and_convolution_invariant(self):
        """Test SDR calculation with scale and conovolution invariant options."""
        estimate = torch.randn(size=(1, 1, 100))
        target = torch.randn(size=(1, 1, 100))

        with pytest.raises(ValueError):
            # using both scale and convolution invariant is not allowed
            calculate_sdr_batch(estimate=estimate, target=target, scale_invariant=True, convolution_invariant=True)

    @pytest.mark.unit
    def test_calculate_mse_input_and_mask(self):
        """Test MSE calculation with simultaneous input length and mask."""
        estimate = torch.randn(size=(1, 1, 100))
        target = torch.randn(size=(1, 1, 100))
        input_length = torch.tensor([100])
        mask = torch.ones(size=(1, 1, 100))

        with pytest.raises(RuntimeError):
            # using both input_length and mask is not allowed
            calculate_mse_batch(estimate=estimate, target=target, input_length=input_length, mask=mask)

    @pytest.mark.unit
    def test_calculate_mse_invalid_dimensions(self):
        """Test MSE calculation with unsupported dimensions."""
        estimate = torch.randn(size=(1, 1, 100, 10))
        target = torch.randn(size=(1, 1, 100))

        with pytest.raises(AssertionError):
            # mismatched dimensions are not allowed
            calculate_mse_batch(estimate=estimate, target=target)

        estimate = torch.randn(size=(1, 1, 100, 10, 20))
        target = torch.randn(size=estimate.shape)

        with pytest.raises(RuntimeError):
            # dimensions larger than four are not allowed
            calculate_mse_batch(estimate=estimate, target=target)

    @pytest.mark.unit
    def test_calculate_mae_input_and_mask(self):
        """Test MAE calculation with simultaneous input length and mask."""
        estimate = torch.randn(size=(1, 1, 100))
        target = torch.randn(size=(1, 1, 100))
        input_length = torch.tensor([100])
        mask = torch.ones(size=(1, 1, 100))

        with pytest.raises(RuntimeError):
            # using both input_length and mask is not allowed
            calculate_mae_batch(estimate=estimate, target=target, input_length=input_length, mask=mask)

    @pytest.mark.unit
    def test_calculate_mae_invalid_dimensions(self):
        """Test MAE calculation with unsupported dimensions."""
        estimate = torch.randn(size=(1, 1, 100, 10))
        target = torch.randn(size=(1, 1, 100))

        with pytest.raises(AssertionError):
            # mismatched dimensions are not allowed
            calculate_mae_batch(estimate=estimate, target=target)

        estimate = torch.randn(size=(1, 1, 100, 10, 20))
        target = torch.randn(size=estimate.shape)

        with pytest.raises(RuntimeError):
            # dimensions larger than four are not allowed
            calculate_mae_batch(estimate=estimate, target=target)

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    def test_sdr(self, num_channels: int):
        """Test SDR calculation"""
        test_eps = [0, 1e-16, 1e-1]
        batch_size = 8
        num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        for remove_mean in [True, False]:
            for eps in test_eps:

                sdr_loss = SDRLoss(eps=eps, remove_mean=remove_mean)

                for n in range(num_batches):

                    # Generate random signal
                    target = _rng.normal(size=(batch_size, num_channels, num_samples))
                    # Random noise + scaling
                    noise = _rng.uniform(low=0.01, high=1) * _rng.normal(size=(batch_size, num_channels, num_samples))
                    # Estimate
                    estimate = target + noise

                    # DC bias for both
                    target += _rng.uniform(low=-1, high=1)
                    estimate += _rng.uniform(low=-1, high=1)

                    # Tensors for testing the loss
                    tensor_estimate = torch.tensor(estimate)
                    tensor_target = torch.tensor(target)

                    # Reference SDR
                    golden_sdr = np.zeros((batch_size, num_channels))
                    for b in range(batch_size):
                        for m in range(num_channels):
                            golden_sdr[b, m] = calculate_sdr_numpy(
                                estimate=estimate[b, m, :],
                                target=target[b, m, :],
                                remove_mean=remove_mean,
                                eps=eps,
                            )

                    # Calculate SDR in torch
                    uut_sdr = calculate_sdr_batch(
                        estimate=tensor_estimate,
                        target=tensor_target,
                        remove_mean=remove_mean,
                        eps=eps,
                    )

                    # Calculate SDR loss
                    uut_sdr_loss = sdr_loss(estimate=tensor_estimate, target=tensor_target)

                    # Compare torch SDR vs numpy
                    assert np.allclose(
                        uut_sdr.cpu().detach().numpy(), golden_sdr, atol=atol
                    ), f'SDR not matching for example {n}, eps={eps}, remove_mean={remove_mean}'

                    # Compare SDR loss vs average of torch SDR
                    assert np.isclose(
                        uut_sdr_loss, -uut_sdr.mean(), atol=atol
                    ), f'SDRLoss not matching for example {n}, eps={eps}, remove_mean={remove_mean}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    def test_sdr_weighted(self, num_channels: int):
        """Test SDR calculation with weighting for channels"""
        batch_size = 8
        num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        channel_weight = _rng.uniform(low=0.01, high=1.0, size=num_channels)
        channel_weight = channel_weight / np.sum(channel_weight)
        sdr_loss = SDRLoss(weight=channel_weight)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=(batch_size, num_channels, num_samples))
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)

            # Reference SDR
            golden_sdr = 0
            for b in range(batch_size):
                sdr = [
                    calculate_sdr_numpy(estimate=estimate[b, m, :], target=target[b, m, :])
                    for m in range(num_channels)
                ]
                # weighted sum
                sdr = np.sum(np.array(sdr) * channel_weight)
                golden_sdr += sdr
            golden_sdr /= batch_size  # average over batch

            # Calculate SDR
            uut_sdr_loss = sdr_loss(estimate=tensor_estimate, target=tensor_target)

            # Compare
            assert np.allclose(
                uut_sdr_loss.cpu().detach().numpy(), -golden_sdr, atol=atol
            ), f'SDRLoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    def test_sdr_input_length(self, num_channels):
        """Test SDR calculation with input length."""
        batch_size = 8
        max_num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        sdr_loss = SDRLoss()

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=(batch_size, num_channels, max_num_samples))
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to random input_length samples
            input_length = _rng.integers(low=1, high=max_num_samples, size=batch_size)

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)
            tensor_input_length = torch.tensor(input_length)

            # Reference SDR
            golden_sdr = 0
            for b, b_len in enumerate(input_length):
                sdr = [
                    calculate_sdr_numpy(estimate=estimate[b, m, :b_len], target=target[b, m, :b_len])
                    for m in range(num_channels)
                ]
                sdr = np.mean(np.array(sdr))
                golden_sdr += sdr
            golden_sdr /= batch_size  # average over batch

            # Calculate SDR
            uut_sdr_loss = sdr_loss(estimate=tensor_estimate, target=tensor_target, input_length=tensor_input_length)

            # Compare
            assert np.allclose(
                uut_sdr_loss.cpu().detach().numpy(), -golden_sdr, atol=atol
            ), f'SDRLoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    def test_sdr_scale_invariant(self, num_channels: int):
        """Test SDR calculation with scale invariant option."""
        batch_size = 8
        max_num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        sdr_loss = SDRLoss(scale_invariant=True)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=(batch_size, num_channels, max_num_samples))
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to random input_length samples
            input_length = _rng.integers(low=1, high=max_num_samples, size=batch_size)

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)
            tensor_input_length = torch.tensor(input_length)

            # Reference SDR
            golden_sdr = 0
            for b, b_len in enumerate(input_length):
                sdr = [
                    calculate_sdr_numpy(
                        estimate=estimate[b, m, :b_len], target=target[b, m, :b_len], scale_invariant=True
                    )
                    for m in range(num_channels)
                ]
                sdr = np.mean(np.array(sdr))
                golden_sdr += sdr
            golden_sdr /= batch_size  # average over batch

            # Calculate SDR loss
            uut_sdr_loss = sdr_loss(estimate=tensor_estimate, target=tensor_target, input_length=tensor_input_length)

            # Compare
            assert np.allclose(
                uut_sdr_loss.cpu().detach().numpy(), -golden_sdr, atol=atol
            ), f'SDRLoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    def test_sdr_binary_mask(self, num_channels):
        """Test SDR calculation with temporal mask."""
        batch_size = 8
        max_num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        sdr_loss = SDRLoss()

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=(batch_size, num_channels, max_num_samples))
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to masked samples
            mask = _rng.integers(low=0, high=2, size=(batch_size, num_channels, max_num_samples))

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)
            tensor_mask = torch.tensor(mask)

            # Reference SDR
            golden_sdr = 0
            for b in range(batch_size):
                sdr = [
                    calculate_sdr_numpy(
                        estimate=estimate[b, m, mask[b, m, :] > 0], target=target[b, m, mask[b, m, :] > 0]
                    )
                    for m in range(num_channels)
                ]
                sdr = np.mean(np.array(sdr))
                golden_sdr += sdr
            golden_sdr /= batch_size  # average over batch

            # Calculate SDR loss
            uut_sdr_loss = sdr_loss(estimate=tensor_estimate, target=tensor_target, mask=tensor_mask)

            # Compare
            assert np.allclose(
                uut_sdr_loss.cpu().detach().numpy(), -golden_sdr, atol=atol
            ), f'SDRLoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1])
    @pytest.mark.parametrize('sdr_max', [10, 0])
    def test_sdr_max(self, num_channels: int, sdr_max: float):
        """Test SDR calculation with soft max threshold."""
        batch_size = 8
        max_num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        sdr_loss = SDRLoss(sdr_max=sdr_max)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=(batch_size, num_channels, max_num_samples))
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to random input_length samples
            input_length = _rng.integers(low=1, high=max_num_samples, size=batch_size)

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)
            tensor_input_length = torch.tensor(input_length)

            # Reference SDR
            golden_sdr = 0
            for b, b_len in enumerate(input_length):
                sdr = [
                    calculate_sdr_numpy(estimate=estimate[b, m, :b_len], target=target[b, m, :b_len], sdr_max=sdr_max)
                    for m in range(num_channels)
                ]
                sdr = np.mean(np.array(sdr))
                golden_sdr += sdr
            golden_sdr /= batch_size  # average over batch

            # Calculate SDR loss
            uut_sdr_loss = sdr_loss(estimate=tensor_estimate, target=tensor_target, input_length=tensor_input_length)

            # Compare
            assert np.allclose(
                uut_sdr_loss.cpu().detach().numpy(), -golden_sdr, atol=atol
            ), f'SDRLoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('filter_length', [1, 32])
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('use_mask', [True, False])
    @pytest.mark.parametrize('use_input_length', [True, False])
    def test_target_calculation(self, num_channels: int, filter_length: int, use_mask: bool, use_input_length: bool):
        """Test target calculation with scale and convolution invariance."""
        batch_size = 8
        max_num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=(batch_size, num_channels, max_num_samples))
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to random input_length samples
            input_length = _rng.integers(low=filter_length, high=max_num_samples, size=batch_size)

            # Corresponding mask
            mask = torch.zeros(size=(batch_size, 1, max_num_samples))
            for i in range(batch_size):
                mask[i, :, : input_length[i]] = 1.0

            if use_mask and use_input_length:
                with pytest.raises(RuntimeError):
                    scale_invariant_target(
                        estimate=torch.tensor(estimate),
                        target=torch.tensor(target),
                        input_length=torch.tensor(input_length),
                        mask=mask,
                    )

                with pytest.raises(RuntimeError):
                    convolution_invariant_target(
                        estimate=torch.tensor(estimate),
                        target=torch.tensor(target),
                        input_length=torch.tensor(input_length),
                        mask=mask,
                        filter_length=filter_length,
                    )

                # Done with this test
                continue

            # UUT
            si_target = scale_invariant_target(
                estimate=torch.tensor(estimate),
                target=torch.tensor(target),
                input_length=torch.tensor(input_length) if use_input_length else None,
                mask=mask if use_mask else None,
            )
            ci_target = convolution_invariant_target(
                estimate=torch.tensor(estimate),
                target=torch.tensor(target),
                input_length=torch.tensor(input_length) if use_input_length else None,
                mask=mask if use_mask else None,
                filter_length=filter_length,
            )

            if filter_length == 1:
                assert torch.allclose(ci_target, si_target), f'SI and CI should match for filter_length=1'

            # Compare against numpy
            for b in range(batch_size):
                # valid length for the current example
                b_len = input_length[b] if use_input_length or use_mask else max_num_samples

                # calculate reference target for each channel
                for m in range(num_channels):
                    # Scale invariant reference
                    si_target_ref = scale_invariant_target_numpy(
                        estimate=estimate[b, m, :b_len], target=target[b, m, :b_len]
                    )

                    assert np.allclose(
                        si_target[b, m, :b_len].cpu().detach().numpy(), si_target_ref, atol=atol
                    ), f'SI not matching for example {n}, channel {m}'

                    # Convolution invariant reference
                    ci_target_ref = convolution_invariant_target_numpy(
                        estimate=estimate[b, m, :b_len], target=target[b, m, :b_len], filter_length=filter_length
                    )

                    assert np.allclose(
                        ci_target[b, m, :b_len].cpu().detach().numpy(), ci_target_ref, atol=atol
                    ), f'CI not matching for example {n}, channel {m}'

    @pytest.mark.unit
    @pytest.mark.parametrize('filter_length', [1, 32])
    @pytest.mark.parametrize('num_channels', [1, 4])
    def test_sdr_convolution_invariant(self, num_channels: int, filter_length: int):
        """Test SDR calculation with convolution invariant option."""
        batch_size = 8
        max_num_samples = 50
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        _rng = np.random.default_rng(seed=random_seed)

        sdr_loss = SDRLoss(convolution_invariant=True, convolution_filter_length=filter_length)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=(batch_size, num_channels, max_num_samples))
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to random input_length samples
            input_length = _rng.integers(low=filter_length, high=max_num_samples, size=batch_size)

            # Calculate SDR loss
            uut_sdr_loss = sdr_loss(
                estimate=torch.tensor(estimate), target=torch.tensor(target), input_length=torch.tensor(input_length)
            )

            # Reference SDR
            golden_sdr = 0
            for b, b_len in enumerate(input_length):
                sdr = [
                    calculate_sdr_numpy(
                        estimate=estimate[b, m, :b_len],
                        target=target[b, m, :b_len],
                        convolution_invariant=True,
                        convolution_filter_length=filter_length,
                    )
                    for m in range(num_channels)
                ]
                sdr = np.mean(np.array(sdr))
                golden_sdr += sdr
            golden_sdr /= batch_size  # average over batch

            # Compare
            assert np.allclose(
                uut_sdr_loss.cpu().detach().numpy(), -golden_sdr, atol=atol
            ), f'SDRLoss not matching for example {n}'

    @pytest.mark.unit
    def test_sdr_scale_and_convolution_invariant(self):
        """Test SDR calculation with scale and conovolution invariant options."""
        with pytest.raises(ValueError):
            # using both scale and convolution invariant is not allowed
            SDRLoss(scale_invariant=True, convolution_invariant=True)

    @pytest.mark.unit
    def test_sdr_length_and_mask(self):
        """Test SDR calculation with simultaneous input length and mask."""

        estimate = torch.randn(size=(1, 1, 100))
        target = torch.randn(size=(1, 1, 100))
        input_length = torch.tensor([100])
        mask = torch.ones(size=(1, 1, 100))

        sdr_loss = SDRLoss(scale_invariant=False, convolution_invariant=False)

        with pytest.raises(RuntimeError):
            # using both input_length and mask is not allowed
            sdr_loss(estimate=estimate, target=target, input_length=input_length, mask=mask)

    @pytest.mark.unit
    def test_sdr_invalid_weight(self):
        """Test SDR with invalid weights."""
        with pytest.raises(ValueError):
            # negative weights are not allowed
            SDRLoss(weight=[-1, 1])

        with pytest.raises(ValueError):
            # weights should sum to 1
            SDRLoss(weight=[0.1, 0.1])

    @pytest.mark.unit
    def test_sdr_invalid_reduction(self):
        """Test SDR with invalid reduction."""
        with pytest.raises(ValueError):
            SDRLoss(reduction='not-mean')

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('ndim', [3, 4])
    def test_mse(self, num_channels: int, ndim: int):
        """Test MSE calculation"""
        batch_size = 8
        num_samples = 50
        num_features = 123
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        signal_shape = (
            (batch_size, num_channels, num_features, num_samples)
            if ndim == 4
            else (batch_size, num_channels, num_samples)
        )

        reduction_dim = (-2, -1) if ndim == 4 else -1

        mse_loss = MSELoss(ndim=ndim)

        _rng = np.random.default_rng(seed=random_seed)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=signal_shape)
            # Random noise + scaling
            noise = _rng.uniform(low=0.01, high=1) * _rng.normal(size=signal_shape)
            # Estimate
            estimate = target + noise

            # DC bias for both
            target += _rng.uniform(low=-1, high=1)
            estimate += _rng.uniform(low=-1, high=1)

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)

            # Reference MSE
            golden_mse = np.zeros((batch_size, num_channels))
            for b in range(batch_size):
                for m in range(num_channels):
                    err = estimate[b, m, :] - target[b, m, :]
                    golden_mse[b, m] = np.mean(np.abs(err) ** 2, axis=reduction_dim)

            # Calculate MSE in torch
            uut_mse = calculate_mse_batch(estimate=tensor_estimate, target=tensor_target)

            # Calculate MSE loss
            uut_mse_loss = mse_loss(estimate=tensor_estimate, target=tensor_target)

            # Compare torch SDR vs numpy
            assert np.allclose(
                uut_mse.cpu().detach().numpy(), golden_mse, atol=atol
            ), f'MSE not matching for example {n}'

            # Compare SDR loss vs average of torch SDR
            assert np.isclose(uut_mse_loss, uut_mse.mean(), atol=atol), f'MSELoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('ndim', [3, 4])
    def test_mse_weighted(self, num_channels: int, ndim: int):
        """Test MSE calculation with weighting for channels"""
        batch_size = 8
        num_samples = 50
        num_features = 123
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        signal_shape = (
            (batch_size, num_channels, num_features, num_samples)
            if ndim == 4
            else (batch_size, num_channels, num_samples)
        )

        reduction_dim = (-2, -1) if ndim == 4 else -1

        _rng = np.random.default_rng(seed=random_seed)

        channel_weight = _rng.uniform(low=0.01, high=1.0, size=num_channels)
        channel_weight = channel_weight / np.sum(channel_weight)
        mse_loss = MSELoss(weight=channel_weight, ndim=ndim)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=signal_shape)
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)

            # Reference MSE
            golden_mse = 0
            for b in range(batch_size):
                mse = [
                    np.mean(np.abs(estimate[b, m, :] - target[b, m, :]) ** 2, axis=reduction_dim)
                    for m in range(num_channels)
                ]
                # weighted sum
                mse = np.sum(np.array(mse) * channel_weight)
                golden_mse += mse
            golden_mse /= batch_size  # average over batch

            # Calculate MSE loss
            uut_mse_loss = mse_loss(estimate=tensor_estimate, target=tensor_target)

            # Compare
            assert np.allclose(
                uut_mse_loss.cpu().detach().numpy(), golden_mse, atol=atol
            ), f'MSELoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('ndim', [3, 4])
    def test_mse_input_length(self, num_channels: int, ndim: int):
        """Test MSE calculation with input length."""
        batch_size = 8
        max_num_samples = 50
        num_features = 123
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        signal_shape = (
            (batch_size, num_channels, num_features, max_num_samples)
            if ndim == 4
            else (batch_size, num_channels, max_num_samples)
        )

        reduction_dim = (-2, -1) if ndim == 4 else -1

        _rng = np.random.default_rng(seed=random_seed)

        mse_loss = MSELoss(ndim=ndim)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=signal_shape)
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to random input_length samples
            input_length = _rng.integers(low=1, high=max_num_samples, size=batch_size)

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)
            tensor_input_length = torch.tensor(input_length)

            # Reference MSE
            golden_mse = 0
            for b, b_len in enumerate(input_length):
                mse = [
                    np.mean(np.abs(estimate[b, m, ..., :b_len] - target[b, m, ..., :b_len]) ** 2, axis=reduction_dim)
                    for m in range(num_channels)
                ]
                mse = np.mean(np.array(mse))
                golden_mse += mse
            golden_mse /= batch_size  # average over batch

            # Calculate MSE
            uut_mse_loss = mse_loss(estimate=tensor_estimate, target=tensor_target, input_length=tensor_input_length)

            # Compare
            assert np.allclose(
                uut_mse_loss.cpu().detach().numpy(), golden_mse, atol=atol
            ), f'MSELoss not matching for example {n}'

    @pytest.mark.unit
    def test_mse_invalid_weight(self):
        """Test MSE with unsupported weights."""
        with pytest.raises(ValueError):
            # negative weights are not allowed
            MSELoss(weight=[-1, 1])

        with pytest.raises(ValueError):
            # weights should sum to 1
            MSELoss(weight=[0.1, 0.1])

    @pytest.mark.unit
    def test_mse_invalid_reduction(self):
        """Test MSE with unsupported reduction."""
        with pytest.raises(ValueError):
            # unsupported reduction
            MSELoss(reduction='not-mean')

    @pytest.mark.unit
    def test_mse_invalid_ndim(self):
        """Test MSE with unsupported dimensions."""
        with pytest.raises(ValueError):
            # supports dimensions 3, 4
            MSELoss(ndim=2)

        with pytest.raises(ValueError):
            # supports dimensions 3, 4
            MSELoss(ndim=5)

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('ndim', [3, 4])
    def test_mae(self, num_channels: int, ndim: int):
        """Test MAE calculation"""
        batch_size = 8
        num_samples = 50
        num_features = 123
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        signal_shape = (
            (batch_size, num_channels, num_features, num_samples)
            if ndim == 4
            else (batch_size, num_channels, num_samples)
        )

        reduction_dim = (-2, -1) if ndim == 4 else -1

        mae_loss = MAELoss(ndim=ndim)

        _rng = np.random.default_rng(seed=random_seed)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=signal_shape)
            # Random noise + scaling
            noise = _rng.uniform(low=0.01, high=1) * _rng.normal(size=signal_shape)
            # Estimate
            estimate = target + noise

            # DC bias for both
            target += _rng.uniform(low=-1, high=1)
            estimate += _rng.uniform(low=-1, high=1)

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)

            # Reference MSE
            golden_mae = np.zeros((batch_size, num_channels))
            for b in range(batch_size):
                for m in range(num_channels):
                    err = estimate[b, m, :] - target[b, m, :]
                    golden_mae[b, m] = np.mean(np.abs(err), axis=reduction_dim)

            # Calculate MSE loss
            uut_mae_loss = mae_loss(estimate=tensor_estimate, target=tensor_target)

            # Compare
            assert np.allclose(
                uut_mae_loss.cpu().detach().numpy(), golden_mae.mean(), atol=atol
            ), f'MAE not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('ndim', [3, 4])
    def test_mae_weighted(self, num_channels: int, ndim: int):
        """Test MAE calculation with weighting for channels"""
        batch_size = 8
        num_samples = 50
        num_features = 123
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        signal_shape = (
            (batch_size, num_channels, num_features, num_samples)
            if ndim == 4
            else (batch_size, num_channels, num_samples)
        )

        reduction_dim = (-2, -1) if ndim == 4 else -1

        _rng = np.random.default_rng(seed=random_seed)

        channel_weight = _rng.uniform(low=0.01, high=1.0, size=num_channels)
        channel_weight = channel_weight / np.sum(channel_weight)
        mae_loss = MAELoss(weight=channel_weight, ndim=ndim)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=signal_shape)
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)

            # Reference MAE
            golden_mae = 0
            for b in range(batch_size):
                mae = [
                    np.mean(np.abs(estimate[b, m, :] - target[b, m, :]), axis=reduction_dim)
                    for m in range(num_channels)
                ]
                # weighted sum
                mae = np.sum(np.array(mae) * channel_weight)
                golden_mae += mae
            golden_mae /= batch_size  # average over batch

            # Calculate MAE loss
            uut_mae_loss = mae_loss(estimate=tensor_estimate, target=tensor_target)

            # Compare
            assert np.allclose(
                uut_mae_loss.cpu().detach().numpy(), golden_mae, atol=atol
            ), f'MAELoss not matching for example {n}'

    @pytest.mark.unit
    @pytest.mark.parametrize('num_channels', [1, 4])
    @pytest.mark.parametrize('ndim', [3, 4])
    def test_mae_input_length(self, num_channels: int, ndim: int):
        """Test MAE calculation with input length."""
        batch_size = 8
        max_num_samples = 50
        num_features = 123
        num_batches = 10
        random_seed = 42
        atol = 1e-6

        signal_shape = (
            (batch_size, num_channels, num_features, max_num_samples)
            if ndim == 4
            else (batch_size, num_channels, max_num_samples)
        )

        reduction_dim = (-2, -1) if ndim == 4 else -1

        _rng = np.random.default_rng(seed=random_seed)

        mae_loss = MAELoss(ndim=ndim)

        for n in range(num_batches):

            # Generate random signal
            target = _rng.normal(size=signal_shape)
            # Random noise + scaling
            noise = _rng.uniform(low=0.001, high=10) * _rng.normal(size=target.shape)
            # Estimate
            estimate = target + noise

            # Limit calculation to random input_length samples
            input_length = _rng.integers(low=1, high=max_num_samples, size=batch_size)

            # Tensors for testing the loss
            tensor_estimate = torch.tensor(estimate)
            tensor_target = torch.tensor(target)
            tensor_input_length = torch.tensor(input_length)

            # Reference MSE
            golden_mae = 0
            for b, b_len in enumerate(input_length):
                mae = [
                    np.mean(np.abs(estimate[b, m, ..., :b_len] - target[b, m, ..., :b_len]), axis=reduction_dim)
                    for m in range(num_channels)
                ]
                mae = np.mean(np.array(mae))
                golden_mae += mae
            golden_mae /= batch_size  # average over batch

            # Calculate MSE
            uut_mae_loss = mae_loss(estimate=tensor_estimate, target=tensor_target, input_length=tensor_input_length)

            # Compare
            assert np.allclose(
                uut_mae_loss.cpu().detach().numpy(), golden_mae, atol=atol
            ), f'MAELoss not matching for example {n}'

    @pytest.mark.unit
    def test_mae_invalid_weight(self):
        """Test MAE with invalid weights."""
        with pytest.raises(ValueError):
            # negative weights are not allowed
            MAELoss(weight=[-1, 1])

        with pytest.raises(ValueError):
            # weights should sum to 1
            MAELoss(weight=[0.1, 0.1])

    @pytest.mark.unit
    def test_mae_invalid_reduction(self):
        """Test MAE with invalid reduction."""
        with pytest.raises(ValueError):
            # unsupported reduction
            MAELoss(reduction='not-mean')

    @pytest.mark.unit
    def test_mae_invalid_ndim(self):
        """Test MAE with invalid dimensions."""
        with pytest.raises(ValueError):
            # supports dimensions 3, 4
            MAELoss(ndim=2)

        with pytest.raises(ValueError):
            # supports dimensions 3, 4
            MAELoss(ndim=5)

    @pytest.mark.unit
    @pytest.mark.skipif(not HAVE_TORCHAUDIO, reason="Modules in this test require torchaudio")
    def test_maxine_combined_loss(self, test_data_dir):
        INPUT_LOCATION = os.path.join(test_data_dir, 'audio', 'maxine', 'input.bin')
        ATOL = 1e-2

        GOLDEN_VALUES = [
            ((1, 0, 0, True, True), 80.0),
            ((0, 1, 0, True, True), 1.9749),
            ((0, 0, 1, True, True), 19.2192),
            ((1, 1, 1, True, True), 101.1941),
        ]
        batch_size = 16

        for value in GOLDEN_VALUES:
            config, golden = value
            sisnr_wt, asr_wt, spec_wt, use_asr, use_spec = config

            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

            loss_instance = CombinedLoss(
                sample_rate=16000,
                fft_length=1920,
                hop_length=480,
                num_mels=320,
                sisnr_loss_weight=sisnr_wt,
                asr_loss_weight=asr_wt,
                spectral_loss_weight=spec_wt,
                use_asr_loss=use_asr,
                use_mel_spec=use_spec,
            ).to(device)

            input_data = torch.tensor(np.fromfile(INPUT_LOCATION, np.float32))
            input_data = input_data.repeat(batch_size).reshape((batch_size, 1, -1))
            input_data = input_data.to(device)

            estimate = torch.tensor(np.zeros(input_data.shape, np.float32)).reshape((batch_size, 1, -1)).to(device)

            loss = loss_instance.forward(estimate=estimate, target=input_data).cpu()

            assert np.allclose(loss, golden, atol=ATOL)