docstr: Simple test to confirm that the group FSQ module can be instantiated and run,
  and that forward produces the same result as encode-decode.
docstr_tok: ["Simple","test","to","confirm","that","the","group","FSQ","module","can","be","instantiated","and","run",",","and","that","forward","produces","the","same","result","as","encode-decode","."]
code: def test_group_fsq_eval(
      self, num_groups: int, num_levels_per_group: int
  ):
      """Simple test to confirm that the group FSQ module can be instantiated and run,
      and that forward produces the same result as encode-decode.
      """
      # Test inference with group FSQ
      # instantiate
      gfsq = GroupFiniteScalarQuantizer(
          num_groups=num_groups,
          num_levels_per_group=num_levels_per_group,
      )
  
      for i in range(self.num_examples):
          inputs = torch.randn(
              [self.batch_size, gfsq.codebook_dim, self.max_len]
          )
          input_len = torch.tensor(
              [self.max_len] * self.batch_size, dtype=torch.int32
          )
  
          # apply forward
          dequantized_fw, indices_fw = gfsq(
              inputs=inputs, input_len=input_len
          )
  
          # encode-decode
          indices_enc = gfsq.encode(inputs=inputs, input_len=input_len)
          dequantized_dec = gfsq.decode(
              indices=indices_enc, input_len=input_len
          )
  
          # make sure the results are the same
          torch.testing.assert_close(
              indices_enc,
              indices_fw,
              msg=f"example {i}: indices mismatch",
          )
          torch.testing.assert_close(
              dequantized_dec,
              dequantized_fw,
              msg=f"example {i}: dequantized mismatch",
          )
  
          # apply individual FSQs and make sure the results are the same
          inputs_grouped = inputs.chunk(num_groups, dim=1)
          dequantized_fw_grouped = dequantized_fw.chunk(
              num_groups, dim=1
          )
          indices_fw_grouped = indices_fw.chunk(num_groups, dim=0)
  
          for g in range(num_groups):
              dequantized, indices = gfsq.fsqs[g](
                  inputs=inputs_grouped[g], input_len=input_len
              )
              torch.testing.assert_close(
                  dequantized,
                  dequantized_fw_grouped[g],
                  msg=f"example {i}: dequantized mismatch for group {g}",
              )
              torch.testing.assert_close(
                  indices,
                  indices_fw_grouped[g],
                  msg=f"example {i}: indices mismatch for group {g}",
              )
  
imports: import pytest
  import torch
  from nemo.collections.tts.modules.audio_codec_modules import CodecActivation, Conv1dNorm, ConvTranspose1dNorm, FiniteScalarQuantizer, GroupFiniteScalarQuantizer, get_down_sample_padding
  from nemo.collections.tts.modules.encodec_modules import GroupResidualVectorQuantizer, ResidualVectorQuantizer
label_desc: Use image node_id 4 for calling the TestFiniteScalarQuantizer obj's underlying member method code with example usage: obj.test_group_fsq_eval(num_groups, num_levels_per_group) without return types