code: def _count_bn(
      module: nn.BatchNorm1d
      | nn.BatchNorm2d
      | nn.BatchNorm3d
      | nas_nn.MutableBatchNorm1d
      | nas_nn.MutableBatchNorm2d
      | nas_nn.MutableBatchNorm3d,
      input: tuple[MutableShape,],
      output: MutableShape,
      count_normalization: bool = True,
  ) -> FlopsResult:
      if not count_normalization:
          return FlopsResult(0.0, 0.0)
      affine = _getattr(module, "affine", expected_type=bool)
      num_features = _getattr(module, "num_features")
  
      x = input[0]
      total_ops = 2 * (1 + affine) * x[1:].numel()
      total_params = affine * 2 * num_features
      return FlopsResult(total_ops, total_params)
  
imports: from __future__ import annotations
  import logging
  from dataclasses import dataclass, fields
  from functools import partial, reduce
  from typing import Any, NamedTuple, Callable, cast
  from torch import nn
  import nni.nas.nn.pytorch
  from nni.mutable import Sample, Mutable, MutableExpression
  from nni.nas.nn.pytorch import ModelSpace
  from nni.nas.profiler import Profiler, ExpressionProfiler
  from .utils import MutableShape, ShapeTensor, submodule_input_output_shapes, standardize_arguments, argument_in_spec, concat_name
  from .utils._attrs import tuple_n_t, _getattr
label_desc: Use image node_id 6 for calling a global function with example usage: _count_bn(module, input, output, count_normalization) and returns: FlopsResult, FlopsResult