code: def _matmul_flop_jit(
      inputs: Tuple[torch.Tensor], outputs: Tuple[Any]
  ) -> Number:
      """
      Count flops for matmul.
      """
      # Inputs should be a list of length 2.
      # Inputs contains the shapes of two matrices.
      input_shapes = [v.shape for v in inputs]
      assert len(input_shapes) == 2, input_shapes
      assert input_shapes[0][-1] == input_shapes[1][-2], input_shapes
      flop = inputs[0].numel() * input_shapes[-1][-1]
      return flop
  
imports: import logging
  import operator
  from collections import defaultdict
  from functools import reduce
  from numbers import Number
  from typing import Any, Callable, DefaultDict, Dict, List, Tuple
  import torch
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils._pytree import PyTree, tree_map
label_desc: Use image node_id 1 for calling a global function with example usage: _matmul_flop_jit(inputs, outputs) and returns: flop