code: def attn_mask(q, k, v, mask):
      sim = einsum("b i d, b j d -> b i j", q, k)
      mask = (1.0 - mask) * -10000.0
      mask = repeat(mask, "b d -> b r d", r=q.shape[1])
      sim = sim + mask
      attn2 = sim.softmax(dim=-1)
      out = einsum("b i j, b j d -> b i d", attn2, v)
      return out
  
imports: from collections import OrderedDict
  from functools import partial
  import torch
  from einops import rearrange, repeat
  from timm.models.layers import DropPath, trunc_normal_
  from torch import einsum, nn
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.frozen_in_time.frozen_video_transformer import attn
label_desc: Use image node_id 1 for calling a global function with example usage: attn_mask(q, k, v, mask) and returns: out