text: # File under the MIT license, see https://github.com/adefossez/julius/LICENSE for details.
  # Author: adefossez, 2020
  
  import random
  import unittest
  
  import torch as th
  from torch.nn import functional as F
  
  import julius
  
  TOLERANCE = 1e-4  # as relative delta in percentage
  
  
  class _BaseTest(unittest.TestCase):
      def setUp(self):
          th.manual_seed(1234)
          random.seed(1234)
  
      def assertSimilar(self, a, b, msg=None, tol=TOLERANCE):
          delta = 100 * th.norm(a - b) / th.norm(b)
          self.assertLessEqual(delta, tol, msg)
  
      def compare_pytorch(self, *args, block_ratio=10, msg=None, tol=TOLERANCE, **kwargs):
          y_ref = F.conv1d(*args, **kwargs)
          y = julius.fft_conv1d(*args, block_ratio=block_ratio, **kwargs)
          self.assertEqual(list(y.shape), list(y_ref.shape), msg)
          self.assertSimilar(y, y_ref, msg, tol)
  
  
  class TestFFTConv1d(_BaseTest):
      def test_same_as_pytorch(self):
          for _ in range(5):
              kernel_size = random.randrange(4, 128)
              batch_size = random.randrange(1, 6)
              length = random.randrange(kernel_size, 1024)
              chin = random.randrange(1, 12)
              chout = random.randrange(1, 12)
              block_ratio = random.choice([5, 10, 20])
              bias = random.random() < 0.5
              if random.random() < 0.5:
                  padding = 0
              else:
                  padding = random.randrange(kernel_size // 2, 2 * kernel_size)
              x = th.randn(batch_size, chin, length)
              w = th.randn(chout, chin, kernel_size)
              keys = ["length", "kernel_size", "chin", "chout", "block_ratio", "bias"]
              loc = dict(locals())
              state = {key: loc[key] for key in keys}
              if bias:
                  bias = th.randn(chout)
              else:
                  bias = None
              for stride in [1, 2, 5]:
                  state["stride"] = stride
                  self.compare_pytorch(
                      x, w, bias, stride, padding, block_ratio=block_ratio,
                      msg=repr(state))
  
      def test_small_input(self):
          x = th.randn(1, 5, 19)
          w = th.randn(10, 5, 32)
          with self.assertRaises(RuntimeError):
              julius.fft_conv1d(x, w)
  
          x = th.randn(1, 5, 19)
          w = th.randn(10, 5, 19)
          self.assertEqual(list(julius.fft_conv1d(x, w).shape), [1, 10, 1])
  
      def test_block_ratio(self):
          x = th.randn(1, 5, 1024)
          w = th.randn(10, 5, 19)
          ref = julius.fft_conv1d(x, w)
          for block_ratio in [1, 5, 10, 20]:
              y = julius.fft_conv1d(x, w, block_ratio=block_ratio)
              self.assertSimilar(y, ref, msg=str(block_ratio))
  
          with self.assertRaises(RuntimeError):
              y = julius.fft_conv1d(x, w, block_ratio=0.9)
  
      def test_module(self):
          x = th.randn(16, 4, 1024)
          mod = julius.FFTConv1d(4, 5, 8, bias=True)
          mod(x)
          mod = julius.FFTConv1d(4, 5, 8, bias=False)
          mod(x)
  
      def test_torchscript(self):
          x = th.randn(16, 4, 1024)
          mod = julius.FFTConv1d(4, 5, 8, bias=True)
          jitted = th.jit.script(mod)
          self.assertEqual(list(jitted(x).shape), [16, 5, 1024 - 8 + 1])
  
      def test_repr(self):
          mod = julius.FFTConv1d(4, 5, 8, bias=False)
          self.assertEqual(
              repr(mod),
              "FFTConv1d(in_channels=4,out_channels=5,kernel_size=8,bias=False)")
  