docstr: Splits hidden_size dim into attn_head_size and num_heads
code: def _split_heads(self, tensor, num_heads, attn_head_size):
      """
      Splits hidden_size dim into attn_head_size and num_heads
      """
      new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
      tensor = tensor.view(new_shape)
      return tensor.permute(0, 2, 1, 3)
  
imports: import os
  from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from ...activations import ACT2FN
  from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
  from ...modeling_utils import PreTrainedModel
  from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
  from .configuration_gpt_neo import GPTNeoConfig
label_desc: Use image node_id 2 for calling the GPTNeoSelfAttention obj's underlying member method code with example usage: obj._split_heads(tensor, num_heads, attn_head_size) and returns: tensor