torch-l1-snr 0.0.3__tar.gz → 0.0.5__tar.gz

{torch_l1_snr-0.0.3/torch_l1_snr.egg-info → torch_l1_snr-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-l1-snr
-Version: 0.0.3
+Version: 0.0.5
 Summary: L1-SNR loss functions for audio source separation in PyTorch
 Home-page: https://github.com/crlandsc/torch-l1-snr
 Author: Christopher Landscaping
@@ -34,13 +34,13 @@ Dynamic: license-file
 
 A PyTorch implementation of L1-based Signal-to-Noise Ratio (SNR) loss functions for audio source separation. This package provides implementations and novel extensions based on concepts from recent academic papers, offering flexible and robust loss functions that can be easily integrated into any PyTorch-based audio separation pipeline.
 
-The core `L1SNRLoss` is based on the loss function described in [1], while `L1SNRDBLoss` and `STFTL1SNRDBLoss` are extensions of the adaptive level-matching regularization technique proposed in [2].
+The core `L1SNRLoss` is based on the loss function described in [[1]](https://arxiv.org/abs/2309.02539), while `L1SNRDBLoss` and `STFTL1SNRDBLoss` are extensions of the adaptive level-matching regularization technique proposed in [[2]](https://arxiv.org/abs/2501.16171).
 
 ## Features
 
-- **Time-Domain L1SNR Loss**: A basic, time-domain L1-SNR loss, based on [1].
-- **Regularized Time-Domain L1SNRDBLoss**: An extension of the L1SNR loss with adaptive level-matching regularization from [2], plus an optional L1 loss component.
-- **Multi-Resolution STFT L1SNRDBLoss**: A spectrogram-domain version of the loss from [2], calculated over multiple STFT resolutions.
+- **Time-Domain L1SNR Loss**: A basic, time-domain L1-SNR loss, based on [[1]](https://arxiv.org/abs/2309.02539).
+- **Regularized Time-Domain L1SNRDBLoss**: An extension of the L1SNR loss with adaptive level-matching regularization from [[2]](https://arxiv.org/abs/2501.16171), plus an optional L1 loss component.
+- **Multi-Resolution STFT L1SNRDBLoss**: A spectrogram-domain version of the loss from [[2]](https://arxiv.org/abs/2501.16171), calculated over multiple STFT resolutions.
 - **Modular Stem-based Loss**: A wrapper that combines time and spectrogram domain losses and can be configured to run on specific stems.
 - **Efficient & Robust**: Includes optimizations for pure L1 loss calculation and robust handling of `NaN`/`inf` values and short audio segments.
 
@@ -166,7 +166,7 @@ The goal of these loss functions is to provide a perceptually-informed and robus
 
 #### Level-Matching Regularization
 
-A key feature of `L1SNRDBLoss` is the adaptive regularization term, as described in [2]. This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (`lambda`) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
+A key feature of `L1SNRDBLoss` is the adaptive regularization term, as described in [[2]](https://arxiv.org/abs/2501.16171). This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (`lambda`) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
 
 #### Multi-Resolution Spectrogram Analysis
 
@@ -205,8 +205,8 @@ The loss functions implemented here are based on the work of the authors of the
 
 ## References
 
-[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, 2023. (arXiv:2309.02539)
+[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, 2023. [arXiv:2309.02539](https://arxiv.org/abs/2309.02539)
 
-[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," arXiv:2501.16171.
+[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," [arXiv:2501.16171](https://arxiv.org/abs/2501.16171).
 
-[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. (arXiv:2406.18747) 
+[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. [arXiv:2406.18747](https://arxiv.org/abs/2406.18747) 

{torch_l1_snr-0.0.3 → torch_l1_snr-0.0.5}/README.md

@@ -6,13 +6,13 @@
 
 A PyTorch implementation of L1-based Signal-to-Noise Ratio (SNR) loss functions for audio source separation. This package provides implementations and novel extensions based on concepts from recent academic papers, offering flexible and robust loss functions that can be easily integrated into any PyTorch-based audio separation pipeline.
 
-The core `L1SNRLoss` is based on the loss function described in [1], while `L1SNRDBLoss` and `STFTL1SNRDBLoss` are extensions of the adaptive level-matching regularization technique proposed in [2].
+The core `L1SNRLoss` is based on the loss function described in [[1]](https://arxiv.org/abs/2309.02539), while `L1SNRDBLoss` and `STFTL1SNRDBLoss` are extensions of the adaptive level-matching regularization technique proposed in [[2]](https://arxiv.org/abs/2501.16171).
 
 ## Features
 
-- **Time-Domain L1SNR Loss**: A basic, time-domain L1-SNR loss, based on [1].
-- **Regularized Time-Domain L1SNRDBLoss**: An extension of the L1SNR loss with adaptive level-matching regularization from [2], plus an optional L1 loss component.
-- **Multi-Resolution STFT L1SNRDBLoss**: A spectrogram-domain version of the loss from [2], calculated over multiple STFT resolutions.
+- **Time-Domain L1SNR Loss**: A basic, time-domain L1-SNR loss, based on [[1]](https://arxiv.org/abs/2309.02539).
+- **Regularized Time-Domain L1SNRDBLoss**: An extension of the L1SNR loss with adaptive level-matching regularization from [[2]](https://arxiv.org/abs/2501.16171), plus an optional L1 loss component.
+- **Multi-Resolution STFT L1SNRDBLoss**: A spectrogram-domain version of the loss from [[2]](https://arxiv.org/abs/2501.16171), calculated over multiple STFT resolutions.
 - **Modular Stem-based Loss**: A wrapper that combines time and spectrogram domain losses and can be configured to run on specific stems.
 - **Efficient & Robust**: Includes optimizations for pure L1 loss calculation and robust handling of `NaN`/`inf` values and short audio segments.
 
@@ -138,7 +138,7 @@ The goal of these loss functions is to provide a perceptually-informed and robus
 
 #### Level-Matching Regularization
 
-A key feature of `L1SNRDBLoss` is the adaptive regularization term, as described in [2]. This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (`lambda`) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
+A key feature of `L1SNRDBLoss` is the adaptive regularization term, as described in [[2]](https://arxiv.org/abs/2501.16171). This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (`lambda`) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
 
 #### Multi-Resolution Spectrogram Analysis
 
@@ -177,8 +177,8 @@ The loss functions implemented here are based on the work of the authors of the
 
 ## References
 
-[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, 2023. (arXiv:2309.02539)
+[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, 2023. [arXiv:2309.02539](https://arxiv.org/abs/2309.02539)
 
-[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," arXiv:2501.16171.
+[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," [arXiv:2501.16171](https://arxiv.org/abs/2501.16171).
 
-[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. (arXiv:2406.18747) 
+[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. [arXiv:2406.18747](https://arxiv.org/abs/2406.18747) 

{torch_l1_snr-0.0.3 → torch_l1_snr-0.0.5}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = torch-l1-snr
-version = 0.0.3
+version = attr: torch_l1snr.__version__
 author = Christopher Landscaping
 author_email = crlandschoot@gmail.com
 description = L1-SNR loss functions for audio source separation in PyTorch

torch_l1_snr-0.0.5/tests/test_losses.py

@@ -0,0 +1,545 @@
+import torch
+import pytest
+from typing import Optional
+from torch_l1snr import (
+    dbrms,
+    L1SNRLoss,
+    L1SNRDBLoss,
+    STFTL1SNRDBLoss,
+    MultiL1SNRDBLoss,
+)
+
+# --- Test Helper: Stem Wrapper ---
+class StemWrappedLoss(torch.nn.Module):
+    """Test helper matching user's pipL1SNRLoss wrapper pattern."""
+    def __init__(self, base_loss, stem_dimension: Optional[int] = None):
+        super().__init__()
+        self.base_loss = base_loss
+        self.stem_dimension = stem_dimension
+    
+    def forward(self, estimates, actuals, *args, **kwargs):
+        if self.stem_dimension is not None:
+            # Handle both [B,S,T] and [B,S,C,T] shapes
+            if estimates.ndim == 3:  # [B, S, T]
+                est_source = estimates[:, self.stem_dimension, :]
+                act_source = actuals[:, self.stem_dimension, :]
+            else:  # [B, S, C, T]
+                est_source = estimates[:, self.stem_dimension, :, :]
+                act_source = actuals[:, self.stem_dimension, :, :]
+            return self.base_loss(est_source, act_source, *args, **kwargs)
+        else:
+            return self.base_loss(estimates, actuals, *args, **kwargs)
+
+# --- Test Fixtures ---
+@pytest.fixture
+def dummy_audio():
+    """Provides a batch of dummy audio signals."""
+    estimates = torch.randn(2, 16000)
+    actuals = torch.randn(2, 16000)
+    # Ensure actuals are not all zero to avoid division by zero in loss
+    actuals[0, :100] += 0.1 
+    return estimates, actuals
+
+@pytest.fixture
+def dummy_stems():
+    """Provides a batch of dummy multi-stem signals."""
+    estimates = torch.randn(2, 4, 1, 16000) # batch, stems, channels, samples
+    actuals = torch.randn(2, 4, 1, 16000)
+    actuals[:, 0, :, :100] += 0.1 # Ensure not all zero
+    return estimates, actuals
+
+@pytest.fixture
+def dummy_stems_3d():
+    """Multi-stem signals: [B, S, T]"""
+    estimates = torch.randn(2, 4, 16000)
+    actuals = torch.randn(2, 4, 16000)
+    actuals[:, 0, :100] += 0.1  # Ensure not all zero
+    return estimates, actuals
+
+@pytest.fixture
+def dummy_stems_4d():
+    """Multi-stem signals: [B, S, C, T]"""
+    estimates = torch.randn(2, 4, 1, 16000)
+    actuals = torch.randn(2, 4, 1, 16000)
+    actuals[:, 0, :, :100] += 0.1
+    return estimates, actuals
+
+# --- Test Functions ---
+
+def test_dbrms():
+    signal = torch.ones(2, 1000) * 0.1
+    # RMS of 0.1 is -20 dB
+    assert torch.allclose(dbrms(signal), torch.tensor([-20.0, -20.0]), atol=1e-4)
+    
+    zeros = torch.zeros(2, 1000)
+    # dbrms of zero should be -80dB with default eps=1e-8
+    assert torch.allclose(dbrms(zeros), torch.tensor([-80.0, -80.0]), atol=1e-4)
+
+def test_l1snr_loss(dummy_audio):
+    estimates, actuals = dummy_audio
+    loss_fn = L1SNRLoss(name="test")
+    loss = loss_fn(estimates, actuals)
+    
+    assert isinstance(loss, torch.Tensor)
+    assert loss.ndim == 0
+    assert not torch.isnan(loss)
+    assert not torch.isinf(loss)
+
+def test_l1snrdb_loss_time(dummy_audio):
+    estimates, actuals = dummy_audio
+    
+    # Test with default settings (L1SNR + Regularization)
+    loss_fn = L1SNRDBLoss(name="test", use_regularization=True, l1_weight=0.0)
+    loss = loss_fn(estimates, actuals)
+    assert loss.ndim == 0 and not torch.isnan(loss)
+
+    # Test without regularization
+    loss_fn_no_reg = L1SNRDBLoss(name="test_no_reg", use_regularization=False, l1_weight=0.0)
+    loss_no_reg = loss_fn_no_reg(estimates, actuals)
+    assert loss_no_reg.ndim == 0 and not torch.isnan(loss_no_reg)
+
+    # Test with L1 loss component
+    loss_fn_l1 = L1SNRDBLoss(name="test_l1", l1_weight=0.2)
+    loss_l1 = loss_fn_l1(estimates, actuals)
+    assert loss_l1.ndim == 0 and not torch.isnan(loss_l1)
+    
+    # Test pure L1 loss mode
+    loss_fn_pure_l1 = L1SNRDBLoss(name="test_pure_l1", l1_weight=1.0)
+    pure_l1_loss = loss_fn_pure_l1(estimates, actuals)
+    # Pure L1 mode uses torch.nn.L1Loss, so compare with manual L1 calculation
+    l1_loss_manual = torch.nn.L1Loss()(
+        estimates.reshape(estimates.shape[0], -1),
+        actuals.reshape(actuals.shape[0], -1)
+    )
+    assert torch.allclose(pure_l1_loss, l1_loss_manual)
+
+def test_stft_l1snrdb_loss(dummy_audio):
+    estimates, actuals = dummy_audio
+    
+    # Test with default settings
+    loss_fn = STFTL1SNRDBLoss(name="test", l1_weight=0.0)
+    loss = loss_fn(estimates, actuals)
+    assert loss.ndim == 0 and not torch.isnan(loss) and not torch.isinf(loss)
+    
+    # Test pure L1 mode
+    loss_fn_pure_l1 = STFTL1SNRDBLoss(name="test_pure_l1", l1_weight=1.0)
+    l1_loss = loss_fn_pure_l1(estimates, actuals)
+    assert l1_loss.ndim == 0 and not torch.isnan(l1_loss) and not torch.isinf(l1_loss)
+
+    # Test with very short audio
+    short_estimates = estimates[:, :500]
+    short_actuals = actuals[:, :500]
+    loss_short = loss_fn(short_estimates, short_actuals)
+    # min_audio_length is 512, so this should fallback to time-domain loss
+    assert loss_short.ndim == 0 and not torch.isnan(loss_short)
+
+def test_stem_multi_loss(dummy_stems):
+    estimates, actuals = dummy_stems
+
+    # Test with a specific stem - users now manage stems manually by slicing
+    # Extract stem 1 (second stem) manually
+    est_stem = estimates[:, 1, ...]  # Shape: [batch, channels, samples]
+    act_stem = actuals[:, 1, ...]
+    loss_fn_stem = MultiL1SNRDBLoss(
+        name="test_loss_stem",
+        spec_weight=0.5,
+        l1_weight=0.1
+    )
+    loss = loss_fn_stem(est_stem, act_stem)
+    assert loss.ndim == 0 and not torch.isnan(loss)
+
+    # Test with all stems jointly - flatten all stems together
+    # Reshape to [batch, -1] to process all stems at once
+    est_all = estimates.reshape(estimates.shape[0], -1)
+    act_all = actuals.reshape(actuals.shape[0], -1)
+    loss_fn_all = MultiL1SNRDBLoss(
+        name="test_loss_all",
+        spec_weight=0.5,
+        l1_weight=0.1
+    )
+    loss_all = loss_fn_all(est_all, act_all)
+    assert loss_all.ndim == 0 and not torch.isnan(loss_all)
+    
+    # Test pure L1 mode on all stems
+    loss_fn_l1 = MultiL1SNRDBLoss(name="l1_only", l1_weight=1.0)
+    l1_loss = loss_fn_l1(est_all, act_all)
+    
+    # Can't easily compute multi-res STFT L1 here, but can check it's not nan
+    assert l1_loss.ndim == 0 and not torch.isnan(l1_loss)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+def test_loss_variants(dummy_audio, l1_weight):
+    """Test L1SNRDBLoss and STFTL1SNRDBLoss with different l1_weights."""
+    estimates, actuals = dummy_audio
+    
+    time_loss_fn = L1SNRDBLoss(name=f"test_time_{l1_weight}", l1_weight=l1_weight)
+    time_loss = time_loss_fn(estimates, actuals)
+    assert not torch.isnan(time_loss) and not torch.isinf(time_loss)
+
+    spec_loss_fn = STFTL1SNRDBLoss(name=f"test_spec_{l1_weight}", l1_weight=l1_weight)
+    spec_loss = spec_loss_fn(estimates, actuals)
+    assert not torch.isnan(spec_loss) and not torch.isinf(spec_loss)
+
+# --- Wrapper-Paradigm Tests ---
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+def test_l1snr_wrapper_all_stems_3d(dummy_stems_3d, l1_weight):
+    """Test L1SNRLoss wrapper with stem_dimension=None on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = L1SNRLoss(name="test", weight=1.0, l1_weight=l1_weight)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+def test_l1snr_wrapper_all_stems_4d(dummy_stems_4d, l1_weight):
+    """Test L1SNRLoss wrapper with stem_dimension=None on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = L1SNRLoss(name="test", weight=1.0, l1_weight=l1_weight)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_l1snr_wrapper_single_stem_3d(dummy_stems_3d, l1_weight, stem_idx):
+    """Test L1SNRLoss wrapper with stem_dimension=k on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = L1SNRLoss(name="test", weight=1.0, l1_weight=l1_weight)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :]
+    act_slice = actuals[:, stem_idx, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_l1snr_wrapper_single_stem_4d(dummy_stems_4d, l1_weight, stem_idx):
+    """Test L1SNRLoss wrapper with stem_dimension=k on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = L1SNRLoss(name="test", weight=1.0, l1_weight=l1_weight)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :, :]
+    act_slice = actuals[:, stem_idx, :, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_reg", [True, False])
+def test_l1snrdb_wrapper_all_stems_3d(dummy_stems_3d, l1_weight, use_reg):
+    """Test L1SNRDBLoss wrapper with stem_dimension=None on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = L1SNRDBLoss(name="test", weight=1.0, l1_weight=l1_weight, use_regularization=use_reg)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_reg", [True, False])
+def test_l1snrdb_wrapper_all_stems_4d(dummy_stems_4d, l1_weight, use_reg):
+    """Test L1SNRDBLoss wrapper with stem_dimension=None on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = L1SNRDBLoss(name="test", weight=1.0, l1_weight=l1_weight, use_regularization=use_reg)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_reg", [True, False])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_l1snrdb_wrapper_single_stem_3d(dummy_stems_3d, l1_weight, use_reg, stem_idx):
+    """Test L1SNRDBLoss wrapper with stem_dimension=k on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = L1SNRDBLoss(name="test", weight=1.0, l1_weight=l1_weight, use_regularization=use_reg)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :]
+    act_slice = actuals[:, stem_idx, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_reg", [True, False])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_l1snrdb_wrapper_single_stem_4d(dummy_stems_4d, l1_weight, use_reg, stem_idx):
+    """Test L1SNRDBLoss wrapper with stem_dimension=k on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = L1SNRDBLoss(name="test", weight=1.0, l1_weight=l1_weight, use_regularization=use_reg)
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :, :]
+    act_slice = actuals[:, stem_idx, :, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+def test_stft_wrapper_all_stems_3d(dummy_stems_3d, l1_weight):
+    """Test STFTL1SNRDBLoss wrapper with stem_dimension=None on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = STFTL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+def test_stft_wrapper_all_stems_4d(dummy_stems_4d, l1_weight):
+    """Test STFTL1SNRDBLoss wrapper with stem_dimension=None on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = STFTL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_stft_wrapper_single_stem_3d(dummy_stems_3d, l1_weight, stem_idx):
+    """Test STFTL1SNRDBLoss wrapper with stem_dimension=k on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = STFTL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :]
+    act_slice = actuals[:, stem_idx, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_stft_wrapper_single_stem_4d(dummy_stems_4d, l1_weight, stem_idx):
+    """Test STFTL1SNRDBLoss wrapper with stem_dimension=k on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = STFTL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :, :]
+    act_slice = actuals[:, stem_idx, :, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+def test_stft_wrapper_short_audio_3d():
+    """Test STFTL1SNRDBLoss wrapper fallback path with short audio [B,S,T]."""
+    estimates = torch.randn(2, 4, 400)  # Short audio
+    actuals = torch.randn(2, 4, 400)
+    actuals[:, 0, :100] += 0.1
+    
+    base_loss = STFTL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=0.0,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=512
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+def test_stft_wrapper_short_audio_4d():
+    """Test STFTL1SNRDBLoss wrapper fallback path with short audio [B,S,C,T]."""
+    estimates = torch.randn(2, 4, 1, 400)  # Short audio
+    actuals = torch.randn(2, 4, 1, 400)
+    actuals[:, 0, :, :100] += 0.1
+    
+    base_loss = STFTL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=0.0,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=512
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_time_reg", [True, False])
+def test_multi_wrapper_all_stems_3d(dummy_stems_3d, l1_weight, use_time_reg):
+    """Test MultiL1SNRDBLoss wrapper with stem_dimension=None on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = MultiL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        use_time_regularization=use_time_reg, use_spec_regularization=False,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_time_reg", [True, False])
+def test_multi_wrapper_all_stems_4d(dummy_stems_4d, l1_weight, use_time_reg):
+    """Test MultiL1SNRDBLoss wrapper with stem_dimension=None on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = MultiL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        use_time_regularization=use_time_reg, use_spec_regularization=False,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_time_reg", [True, False])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_multi_wrapper_single_stem_3d(dummy_stems_3d, l1_weight, use_time_reg, stem_idx):
+    """Test MultiL1SNRDBLoss wrapper with stem_dimension=k on [B,S,T]."""
+    estimates, actuals = dummy_stems_3d
+    base_loss = MultiL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        use_time_regularization=use_time_reg, use_spec_regularization=False,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :]
+    act_slice = actuals[:, stem_idx, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("use_time_reg", [True, False])
+@pytest.mark.parametrize("stem_idx", [0, 3])
+def test_multi_wrapper_single_stem_4d(dummy_stems_4d, l1_weight, use_time_reg, stem_idx):
+    """Test MultiL1SNRDBLoss wrapper with stem_dimension=k on [B,S,C,T]."""
+    estimates, actuals = dummy_stems_4d
+    base_loss = MultiL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=l1_weight,
+        use_time_regularization=use_time_reg, use_spec_regularization=False,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=256
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=stem_idx)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    est_slice = estimates[:, stem_idx, :, :]
+    act_slice = actuals[:, stem_idx, :, :]
+    direct_result = base_loss(est_slice, act_slice)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+def test_multi_wrapper_short_audio_3d():
+    """Test MultiL1SNRDBLoss wrapper fallback path with short audio [B,S,T]."""
+    estimates = torch.randn(2, 4, 400)  # Short audio
+    actuals = torch.randn(2, 4, 400)
+    actuals[:, 0, :100] += 0.1
+    
+    base_loss = MultiL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=0.0,
+        use_time_regularization=True, use_spec_regularization=False,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=512
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result)
+
+def test_multi_wrapper_short_audio_4d():
+    """Test MultiL1SNRDBLoss wrapper fallback path with short audio [B,S,C,T]."""
+    estimates = torch.randn(2, 4, 1, 400)  # Short audio
+    actuals = torch.randn(2, 4, 1, 400)
+    actuals[:, 0, :, :100] += 0.1
+    
+    base_loss = MultiL1SNRDBLoss(
+        name="test", weight=1.0, l1_weight=0.0,
+        use_time_regularization=True, use_spec_regularization=False,
+        n_ffts=[256], hop_lengths=[64], win_lengths=[256], min_audio_length=512
+    )
+    wrapped_loss = StemWrappedLoss(base_loss, stem_dimension=None)
+    
+    wrapped_result = wrapped_loss(estimates, actuals)
+    direct_result = base_loss(estimates, actuals)
+    
+    assert torch.allclose(wrapped_result, direct_result, atol=1e-6)
+    assert wrapped_result.ndim == 0
+    assert not torch.isnan(wrapped_result) and not torch.isinf(wrapped_result) 

{torch_l1_snr-0.0.3 → torch_l1_snr-0.0.5/torch_l1_snr.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch-l1-snr
-Version: 0.0.3
+Version: 0.0.5
 Summary: L1-SNR loss functions for audio source separation in PyTorch
 Home-page: https://github.com/crlandsc/torch-l1-snr
 Author: Christopher Landscaping
@@ -34,13 +34,13 @@ Dynamic: license-file
 
 A PyTorch implementation of L1-based Signal-to-Noise Ratio (SNR) loss functions for audio source separation. This package provides implementations and novel extensions based on concepts from recent academic papers, offering flexible and robust loss functions that can be easily integrated into any PyTorch-based audio separation pipeline.
 
-The core `L1SNRLoss` is based on the loss function described in [1], while `L1SNRDBLoss` and `STFTL1SNRDBLoss` are extensions of the adaptive level-matching regularization technique proposed in [2].
+The core `L1SNRLoss` is based on the loss function described in [[1]](https://arxiv.org/abs/2309.02539), while `L1SNRDBLoss` and `STFTL1SNRDBLoss` are extensions of the adaptive level-matching regularization technique proposed in [[2]](https://arxiv.org/abs/2501.16171).
 
 ## Features
 
-- **Time-Domain L1SNR Loss**: A basic, time-domain L1-SNR loss, based on [1].
-- **Regularized Time-Domain L1SNRDBLoss**: An extension of the L1SNR loss with adaptive level-matching regularization from [2], plus an optional L1 loss component.
-- **Multi-Resolution STFT L1SNRDBLoss**: A spectrogram-domain version of the loss from [2], calculated over multiple STFT resolutions.
+- **Time-Domain L1SNR Loss**: A basic, time-domain L1-SNR loss, based on [[1]](https://arxiv.org/abs/2309.02539).
+- **Regularized Time-Domain L1SNRDBLoss**: An extension of the L1SNR loss with adaptive level-matching regularization from [[2]](https://arxiv.org/abs/2501.16171), plus an optional L1 loss component.
+- **Multi-Resolution STFT L1SNRDBLoss**: A spectrogram-domain version of the loss from [[2]](https://arxiv.org/abs/2501.16171), calculated over multiple STFT resolutions.
 - **Modular Stem-based Loss**: A wrapper that combines time and spectrogram domain losses and can be configured to run on specific stems.
 - **Efficient & Robust**: Includes optimizations for pure L1 loss calculation and robust handling of `NaN`/`inf` values and short audio segments.
 
@@ -166,7 +166,7 @@ The goal of these loss functions is to provide a perceptually-informed and robus
 
 #### Level-Matching Regularization
 
-A key feature of `L1SNRDBLoss` is the adaptive regularization term, as described in [2]. This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (`lambda`) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
+A key feature of `L1SNRDBLoss` is the adaptive regularization term, as described in [[2]](https://arxiv.org/abs/2501.16171). This component calculates the difference in decibel-scaled root-mean-square (dBRMS) levels between the estimated and actual signals. An adaptive weight (`lambda`) is applied to this difference, which increases when the model incorrectly silences a non-silent target. This encourages the model to learn the correct output level and specifically avoids the model collapsing to a trivial silent solution when uncertain.
 
 #### Multi-Resolution Spectrogram Analysis
 
@@ -205,8 +205,8 @@ The loss functions implemented here are based on the work of the authors of the
 
 ## References
 
-[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, 2023. (arXiv:2309.02539)
+[1] K. N. Watcharasupat, C.-W. Wu, Y. Ding, I. Orife, A. J. Hipple, P. A. Williams, S. Kramer, A. Lerch, and W. Wolcott, "A Generalized Bandsplit Neural Network for Cinematic Audio Source Separation," IEEE Open Journal of Signal Processing, 2023. [arXiv:2309.02539](https://arxiv.org/abs/2309.02539)
 
-[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," arXiv:2501.16171.
+[2] K. N. Watcharasupat and A. Lerch, "Separate This, and All of these Things Around It: Music Source Separation via Hyperellipsoidal Queries," [arXiv:2501.16171](https://arxiv.org/abs/2501.16171).
 
-[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. (arXiv:2406.18747) 
+[3] K. N. Watcharasupat and A. Lerch, "A Stem-Agnostic Single-Decoder System for Music Source Separation Beyond Four Stems," Proceedings of the 25th International Society for Music Information Retrieval Conference, 2024. [arXiv:2406.18747](https://arxiv.org/abs/2406.18747) 

{torch_l1_snr-0.0.3 → torch_l1_snr-0.0.5}/torch_l1snr/__init__.py

@@ -12,4 +12,6 @@ __all__ = [
     "L1SNRDBLoss",
     "STFTL1SNRDBLoss",
     "MultiL1SNRDBLoss",
-] 
+] 
+
+__version__ = "0.0.5"

torch_l1_snr-0.0.3/tests/test_losses.py

@@ -1,143 +0,0 @@
-import torch
-import pytest
-from torch_l1snr import (
-    dbrms,
-    L1SNRLoss,
-    L1SNRDBLoss,
-    STFTL1SNRDBLoss,
-    MultiL1SNRDBLoss,
-)
-
-# --- Test Fixtures ---
-@pytest.fixture
-def dummy_audio():
-    """Provides a batch of dummy audio signals."""
-    estimates = torch.randn(2, 16000)
-    actuals = torch.randn(2, 16000)
-    # Ensure actuals are not all zero to avoid division by zero in loss
-    actuals[0, :100] += 0.1 
-    return estimates, actuals
-
-@pytest.fixture
-def dummy_stems():
-    """Provides a batch of dummy multi-stem signals."""
-    estimates = torch.randn(2, 4, 1, 16000) # batch, stems, channels, samples
-    actuals = torch.randn(2, 4, 1, 16000)
-    actuals[:, 0, :, :100] += 0.1 # Ensure not all zero
-    return estimates, actuals
-
-# --- Test Functions ---
-
-def test_dbrms():
-    signal = torch.ones(2, 1000) * 0.1
-    # RMS of 0.1 is -20 dB
-    assert torch.allclose(dbrms(signal), torch.tensor([-20.0, -20.0]), atol=1e-4)
-    
-    zeros = torch.zeros(2, 1000)
-    # dbrms of zero should be -80dB with default eps=1e-8
-    assert torch.allclose(dbrms(zeros), torch.tensor([-80.0, -80.0]), atol=1e-4)
-
-def test_l1snr_loss(dummy_audio):
-    estimates, actuals = dummy_audio
-    loss_fn = L1SNRLoss(name="test")
-    loss = loss_fn(estimates, actuals)
-    
-    assert isinstance(loss, torch.Tensor)
-    assert loss.ndim == 0
-    assert not torch.isnan(loss)
-    assert not torch.isinf(loss)
-
-def test_l1snrdb_loss_time(dummy_audio):
-    estimates, actuals = dummy_audio
-    
-    # Test with default settings (L1SNR + Regularization)
-    loss_fn = L1SNRDBLoss(name="test", use_regularization=True, l1_weight=0.0)
-    loss = loss_fn(estimates, actuals)
-    assert loss.ndim == 0 and not torch.isnan(loss)
-
-    # Test without regularization
-    loss_fn_no_reg = L1SNRDBLoss(name="test_no_reg", use_regularization=False, l1_weight=0.0)
-    loss_no_reg = loss_fn_no_reg(estimates, actuals)
-    assert loss_no_reg.ndim == 0 and not torch.isnan(loss_no_reg)
-
-    # Test with L1 loss component
-    loss_fn_l1 = L1SNRDBLoss(name="test_l1", l1_weight=0.2)
-    loss_l1 = loss_fn_l1(estimates, actuals)
-    assert loss_l1.ndim == 0 and not torch.isnan(loss_l1)
-    
-    # Test pure L1 loss mode
-    loss_fn_pure_l1 = L1SNRDBLoss(name="test_pure_l1", l1_weight=1.0)
-    pure_l1_loss = loss_fn_pure_l1(estimates, actuals)
-    # Pure L1 mode uses torch.nn.L1Loss, so compare with manual L1 calculation
-    l1_loss_manual = torch.nn.L1Loss()(
-        estimates.reshape(estimates.shape[0], -1),
-        actuals.reshape(actuals.shape[0], -1)
-    )
-    assert torch.allclose(pure_l1_loss, l1_loss_manual)
-
-def test_stft_l1snrdb_loss(dummy_audio):
-    estimates, actuals = dummy_audio
-    
-    # Test with default settings
-    loss_fn = STFTL1SNRDBLoss(name="test", l1_weight=0.0)
-    loss = loss_fn(estimates, actuals)
-    assert loss.ndim == 0 and not torch.isnan(loss) and not torch.isinf(loss)
-    
-    # Test pure L1 mode
-    loss_fn_pure_l1 = STFTL1SNRDBLoss(name="test_pure_l1", l1_weight=1.0)
-    l1_loss = loss_fn_pure_l1(estimates, actuals)
-    assert l1_loss.ndim == 0 and not torch.isnan(l1_loss) and not torch.isinf(l1_loss)
-
-    # Test with very short audio
-    short_estimates = estimates[:, :500]
-    short_actuals = actuals[:, :500]
-    loss_short = loss_fn(short_estimates, short_actuals)
-    # min_audio_length is 512, so this should fallback to time-domain loss
-    assert loss_short.ndim == 0 and not torch.isnan(loss_short)
-
-def test_stem_multi_loss(dummy_stems):
-    estimates, actuals = dummy_stems
-
-    # Test with a specific stem - users now manage stems manually by slicing
-    # Extract stem 1 (second stem) manually
-    est_stem = estimates[:, 1, ...]  # Shape: [batch, channels, samples]
-    act_stem = actuals[:, 1, ...]
-    loss_fn_stem = MultiL1SNRDBLoss(
-        name="test_loss_stem",
-        spec_weight=0.5,
-        l1_weight=0.1
-    )
-    loss = loss_fn_stem(est_stem, act_stem)
-    assert loss.ndim == 0 and not torch.isnan(loss)
-
-    # Test with all stems jointly - flatten all stems together
-    # Reshape to [batch, -1] to process all stems at once
-    est_all = estimates.reshape(estimates.shape[0], -1)
-    act_all = actuals.reshape(actuals.shape[0], -1)
-    loss_fn_all = MultiL1SNRDBLoss(
-        name="test_loss_all",
-        spec_weight=0.5,
-        l1_weight=0.1
-    )
-    loss_all = loss_fn_all(est_all, act_all)
-    assert loss_all.ndim == 0 and not torch.isnan(loss_all)
-    
-    # Test pure L1 mode on all stems
-    loss_fn_l1 = MultiL1SNRDBLoss(name="l1_only", l1_weight=1.0)
-    l1_loss = loss_fn_l1(est_all, act_all)
-    
-    # Can't easily compute multi-res STFT L1 here, but can check it's not nan
-    assert l1_loss.ndim == 0 and not torch.isnan(l1_loss)
-
-@pytest.mark.parametrize("l1_weight", [0.0, 0.5, 1.0])
-def test_loss_variants(dummy_audio, l1_weight):
-    """Test L1SNRDBLoss and STFTL1SNRDBLoss with different l1_weights."""
-    estimates, actuals = dummy_audio
-    
-    time_loss_fn = L1SNRDBLoss(name=f"test_time_{l1_weight}", l1_weight=l1_weight)
-    time_loss = time_loss_fn(estimates, actuals)
-    assert not torch.isnan(time_loss) and not torch.isinf(time_loss)
-
-    spec_loss_fn = STFTL1SNRDBLoss(name=f"test_spec_{l1_weight}", l1_weight=l1_weight)
-    spec_loss = spec_loss_fn(estimates, actuals)
-    assert not torch.isnan(spec_loss) and not torch.isinf(spec_loss)