code: def _calculate_scale(
      head_dim_size: int, scale: Optional[float]
  ) -> float:
      """
      For FlashAttention we pad the head dimension to be a multiple of 8 so we need to scale the output
      by the original head size and not the padded.
      """
      if scale is not None:
          return scale
      return 1.0 / math.sqrt(head_dim_size)
  
imports: import math
  from typing import List, Optional
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: _calculate_scale(head_dim_size, scale) and returns: unknown, scale