code: def test_flop_count_elementwise(input_shape):
      # Tests that the flop counter counts flops correctly (i.e.
      # when compared to a reference value) for a model involving
      # lots of binary and elementwise operations.
  
      class ElementwiseModel(torch.nn.Module):
          def forward(self, inp):
              flops = 5  # view operations
              out = inp
  
              out = out * 2
              # +1 here and below for counting the additions to the flop
              # variable itself.
              flops += out.numel() + 1
  
              out = out.pow(2).exp().log()
              flops += out.numel() * 3 + 1
  
              out = out.view(-1, 1) @ out.view(1, -1)
              flops += out.numel() + 1
  
              out = out + out
              flops += out.numel() + 1
  
              return out, flops
  
      model = ElementwiseModel()
      inp = torch.randn(input_shape)
      _, correct_flops = model(inp)
  
      counted_flops, _, unsupported_ops = count_flops(
          model, (inp,), use_single_flop_estimates=True
      )
      assert counted_flops == correct_flops
      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 3 for calling a global function with example usage: test_flop_count_elementwise(input_shape) and returns: out,  flops