code: def fill_with_head_ids(head_size: int, n_heads: int) -> torch.Tensor:
      """
      Fills a tensor with the associated head ids. All columns should have the same value.
      """
      head_ids = torch.arange(
          n_heads,
          dtype=torch.half,
          device=get_accelerator().current_device(),
      )
  
      head_ids = (
          head_ids.repeat_interleave(head_size)
          .repeat(head_size * n_heads)
          .reshape(n_heads * head_size, -1)
      )
      return head_ids
  
file_path: deepspeed/tests/unit/inference/v2/model_implementations/sharding/test_attn_out_sharding.py
imports: 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) and returns: head_ids