code: def sdpa_flop(
      query_shape,
      key_shape,
      value_shape,
      *args,
      out_shape=None,
      **kwargs,
  ) -> int:
      """Count flops for self-attention."""
      # NB: We aren't accounting for causal attention here
      return sdpa_flop_count(query_shape, key_shape, value_shape)
  
imports: import torch
  import torch.nn
  from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
  from typing import List, Any, Dict, Optional, Union, NamedTuple
  from collections import defaultdict
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils.hooks import RemovableHandle
  from torch._decomp import register_decomposition
  from math import prod
  from functools import wraps
label_desc: Use image node_id 13 for calling a global function with example usage: sdpa_flop(query_shape, key_shape, value_shape) and returns: sdpa_flop_count