docstr: Make sure SpecAug masking works as expected
docstr_tok: ["Make","sure","SpecAug","masking","works","as","expected"]
code: def test_specaugment(
      self,
      n_time_masks,
      time_mask_param,
      n_freq_masks,
      freq_mask_param,
      iid_masks,
  ):
      """Make sure SpecAug masking works as expected"""
      spec = torch.ones(2, 200, 200)
      transform = T.SpecAugment(
          n_time_masks=n_time_masks,
          time_mask_param=time_mask_param,
          n_freq_masks=n_freq_masks,
          freq_mask_param=freq_mask_param,
          iid_masks=iid_masks,
          zero_masking=True,
      )
      spec_masked = transform(spec)
      f_axis_mean = torch.mean(spec_masked, 1)
      t_axis_mean = torch.mean(spec_masked, 2)
      if n_time_masks == 0 and n_freq_masks == 0:
          self.assertEqual(spec, spec_masked)
      elif n_time_masks > 0 and n_freq_masks > 0:
          # Across both time and frequency dimensions, the mean tensor should contain
          # at least one zero element, and all non-zero elements should be less than 1.
          self.assertTrue(0 in t_axis_mean)
          self.assertFalse(
              False in torch.lt(t_axis_mean[t_axis_mean != 0], 1)
          )
          self.assertTrue(0 in f_axis_mean)
          self.assertFalse(
              False in torch.lt(f_axis_mean[f_axis_mean != 0], 1)
          )
      elif n_freq_masks > 0:
          # Across the frequency axis where we apply masking,
          # the mean tensor should contain equal elements,
          # and the value should be between 0 and 1.
          self.assertFalse(
              False in torch.eq(f_axis_mean[0], f_axis_mean[0][0])
          )
          self.assertFalse(
              False in torch.eq(f_axis_mean[1], f_axis_mean[1][0])
          )
          self.assertTrue(f_axis_mean[0][0] < 1)
          self.assertTrue(f_axis_mean[1][0] > 0)
  
          # Across the time axis where we don't mask, the mean tensor should contain at
          # least one zero element, and all non-zero elements should be 1.
          self.assertTrue(0 in t_axis_mean)
          self.assertFalse(
              False in torch.eq(t_axis_mean[t_axis_mean != 0], 1)
          )
      else:
          # Across the time axis where we apply masking,
          # the mean tensor should contain equal elements,
          # and the value should be between 0 and 1.
          self.assertFalse(
              False in torch.eq(t_axis_mean[0], t_axis_mean[0][0])
          )
          self.assertFalse(
              False in torch.eq(t_axis_mean[1], t_axis_mean[1][0])
          )
          self.assertTrue(t_axis_mean[0][0] < 1)
          self.assertTrue(t_axis_mean[1][0] > 0)
  
          # Across the frequency axis where we don't mask, the mean tensor should contain at
          # least one zero element, and all non-zero elements should be 1.
          self.assertTrue(0 in f_axis_mean)
          self.assertFalse(
              False in torch.eq(f_axis_mean[f_axis_mean != 0], 1)
          )
  
      # Test if iid_masks gives different masking results for different spectrograms across the 0th dimension.
      diff = torch.linalg.vector_norm(
          spec_masked[0] - spec_masked[1]
      ).item()
      print(diff)
      if iid_masks is True:
          self.assertTrue(diff > 0)
      else:
          self.assertTrue(diff == 0)
  
imports: import math
  import random
  from unittest.mock import patch
  import numpy
  import torch
  import torchaudio.transforms
  from parameterized import param, parameterized
  from scipy import signal
  from torchaudio.functional import lfilter, preemphasis
  from torchaudio.functional.functional import _get_sinc_resample_kernel
  from torchaudio_unittest.common_utils import get_spectrogram, get_whitenoise, nested_params, TestBaseMixin
  from torchaudio_unittest.common_utils.psd_utils import psd_numpy
label_desc: Use image node_id 20 for calling the TransformsTestBase obj's underlying member method code with example usage: obj.test_specaugment(n_time_masks, time_mask_param, n_freq_masks, freq_mask_param, iid_masks) without return types