code: def build_attention_mask(
      context_length: int, batch_size: int, use_pytorch_mha
  ):
      # lazily create causal attention mask, with full attention between the vision tokens
      # pytorch uses additive attention mask; fill with -inf
      mask = torch.empty(context_length, context_length)
      mask.fill_(float("-inf"))
      mask.triu_(1)  # zero out the lower diagonal
      if not use_pytorch_mha:
          mask = mask.unsqueeze(0)  # add dummy batch dimension
          mask = mask.expand(batch_size, -1, -1)
      return mask
  
imports: import sys
  import numpy
  import pytest
  import torch
  from cvnets.layers.multi_head_attention import MultiHeadAttention
label_desc: Use image node_id 1 for calling a global function with example usage: build_attention_mask(context_length, batch_size, use_pytorch_mha) and returns: mask