code: def template(self) -> str:
      if self.direction == "fwd":
          return KERNEL_IMPL_TEMPLATE_FWD.format(
              DTYPE=DTYPE_MAP[self.dtype], HEAD_DIM=self.head_dim
          )
      elif self.direction == "bwd":
          return KERNEL_IMPL_TEMPLATE_BWD.format(
              DTYPE=DTYPE_MAP[self.dtype], HEAD_DIM=self.head_dim
          )
      else:
          return KERNEL_IMPL_TEMPLATE_FWD_SPLIT.format(
              DTYPE=DTYPE_MAP[self.dtype], HEAD_DIM=self.head_dim
          )
  
file_path: pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/kernels/generate_kernels.py
imports: import argparse
  import itertools
  from dataclasses import dataclass
  from pathlib import Path
  from typing import List, Optional
label_desc: Use image node_id 1 for calling the Kernel obj's underlying member method code with example usage: obj.template() and returns: KERNEL_IMPL_TEMPLATE_FWD, KERNEL_IMPL_TEMPLATE_BWD, KERNEL_IMPL_TEMPLATE_FWD_SPLIT