code: def _get_tensor_bias_bmk(
      attn_bias: Optional[Union[torch.Tensor, AttentionBias]]
  ) -> Optional[torch.Tensor]:
      if not isinstance(attn_bias, torch.Tensor):
          assert attn_bias is None
          return None
      # BMK -> BMHK
      if attn_bias.ndim == 4:
          attn_bias = attn_bias.reshape([-1, *attn_bias.shape[2:]])
      return attn_bias
  
file_path: xformers/xformers/ops/fmha/small_k.py
imports: from typing import Any, List, Mapping, Optional, Set, Tuple, Union
  import torch
  from ..common import get_xformers_operator, register_operator
  from .attn_bias import AttentionBias
  from .common import AttentionBwOpBase, AttentionFwOpBase, Context, Gradients, Inputs, bmk2bmhk
label_desc: Use image node_id 2 for calling a global function with example usage: _get_tensor_bias_bmk(attn_bias) and returns: attn_bias, None