code: def count_flops_attn(model, _x, y):
      """
      A counter for the `thop` package to count the operations in an
      attention operation.
      Meant to be used like:
          macs, params = thop.profile(
              model,
              inputs=(inputs, timestamps),
              custom_ops={QKVAttention: QKVAttention.count_flops},
          )
      """
      b, c, *spatial = y[0].shape
      num_spatial = int(np.prod(spatial))
      # We perform two matmuls with the same number of ops.
      # The first computes the weight matrix, the second computes
      # the combination of the value vectors.
      matmul_ops = 2 * b * (num_spatial**2) * c
      model.total_ops += paddle.to_tensor([matmul_ops], dtype="float64")
  
imports: import math
  from abc import abstractmethod
  import numpy
  import paddle
  import paddle.nn
  import paddle.nn.functional
  from .nn import avg_pool_nd
  from .nn import checkpoint
  from .nn import conv_nd
  from .nn import linear
  from .nn import normalization
  from .nn import SiLU
  from .nn import timestep_embedding
  from .nn import zero_module
label_desc: Use image node_id 1 for calling a global function with example usage: count_flops_attn(model, _x, y) without return types