code: def test_flop_count_ffn(batch_size, dim_in, dim_out):
      # Testing the flop counter counts simple modules as expected.
  
      ffn_model = torch.nn.Sequential(
          torch.nn.Linear(dim_in, dim_out, bias=False),
          torch.nn.Linear(dim_out, dim_out, bias=False),
      )
      ffn_inputs = (torch.randn(batch_size, dim_in),)
  
      total_flops, _, unsupported_ops = count_flops(
          ffn_model, ffn_inputs
      )
      assert (
          total_flops
          == batch_size * dim_in * dim_out
          + batch_size * dim_out * dim_out
      )
      assert len(unsupported_ops) == 0
  
imports: import pytest
  import torch
  import torch.nn
  from neuralcompression.functional import count_flops
label_desc: Use image node_id 1 for calling a global function with example usage: test_flop_count_ffn(batch_size, dim_in, dim_out) without return types