docstr: Splits hidden dim into attn_head_size and num_attention_heads
docstr_tok: ["Splits","hidden","dim","into","attn_head_size","and","num_attention_heads"]
code: def _split_heads(cls, tensor, num_attention_heads, attn_head_size):
      """
      Splits hidden dim into attn_head_size and num_attention_heads
      """
      # tensor: [bs, seq_len, hidden_size]
      new_shape = tensor.size()[:-1] + (
          num_attention_heads,
          attn_head_size,
      )
      # -> [bs, seq_len, num_attention_heads, attn_head_size]
      tensor = tensor.view(new_shape)
      # -> [bs, num_attention_heads, seq_len, attn_head_size]
      tensor = tensor.permute(0, 2, 1, 3)
      return tensor
  
imports: from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import Tensor, nn
  from torch.nn import CrossEntropyLoss
  from ...activations import ACT2FN
  from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
  from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from ...modeling_utils import PreTrainedModel
  from ...utils import logging
  from .configuration_gpt_neox_japanese import GPTNeoXJapaneseConfig
label_desc: Use image node_id 3 for calling the GPTNeoXJapaneseAttention obj's underlying member method code with example usage: obj._split_heads(cls, tensor, num_attention_heads, attn_head_size) and returns: tensor