code: def fill_with_head_ids(
      head_size: int, n_heads_q: int, n_heads_kv: Optional[int] = None
  ) -> torch.Tensor:
      """ """
      head_ids_q = torch.arange(
          n_heads_q,
          dtype=torch.half,
          device=get_accelerator().current_device(),
      )
      head_vals_q = head_ids_q.repeat_interleave(
          head_size * head_size * n_heads_q
      ).reshape(n_heads_q * head_size, -1)
  
      if n_heads_kv is None:
          return torch.cat(
              [head_vals_q, head_vals_q, head_vals_q], dim=0
          )
  
      head_ids_k = torch.arange(
          n_heads_kv,
          dtype=torch.half,
          device=get_accelerator().current_device(),
      )
      head_vals_k = head_ids_k.repeat_interleave(
          head_size * head_size * n_heads_q
      ).reshape(n_heads_kv * head_size, -1)
  
      return torch.cat([head_vals_q, head_vals_k, head_vals_k], dim=0)
  
file_path: DeepSpeed/tests/unit/inference/v2/model_implementations/sharding/test_qkv_sharding.py
imports: from typing import Optional
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.model_implementations.sharding import *
label_desc: Use image node_id 1 for calling a global function with example usage: fill_with_head_ids(head_size, n_heads_q, n_heads_kv) and returns: torch, torch