code: def norm_flop_counter(
      affine_arg_index: int, input_arg_index: int
  ) -> Callable:
      """
      Args:
          affine_arg_index: index of the affine argument in inputs
      """
  
      def norm_flop_jit(
          inputs: List[Any], outputs: List[Any]
      ) -> Number:
          """
          Count flops for norm layers.
          """
          # Inputs[0] contains the shape of the input.
          input_shape = inputs[input_arg_index].shape
  
          has_affine = (
              inputs[affine_arg_index].shape is not None
              if hasattr(inputs[affine_arg_index], "shape")
              else inputs[affine_arg_index]
          )
          assert 2 <= len(input_shape) <= 5, input_shape
          # 5 is just a rough estimate
          flop = reduce(operator.mul, input_shape) * (
              5 if has_affine else 4
          )
          return flop
  
      return norm_flop_jit
  
imports: import operator
  from collections import defaultdict
  from contextlib import contextmanager
  from enum import Enum, auto
  from functools import partial, reduce
  from numbers import Number
  from typing import Any, Callable, List, Union
  import torch
  from packaging import version
  from torch.utils._pytree import tree_map
  from .meta_tensor import MetaTensor
label_desc: Use image node_id 12 for calling a global function with example usage: norm_flop_counter(affine_arg_index, input_arg_index) and returns: norm_flop_jit, flop