docstr: Splits hidden dim into attn_head_size and num_attention_heads
code: def _split_heads(
      self, tensor, num_attention_heads, attn_head_size, rotary
  ):
      """
      Splits hidden dim into attn_head_size and num_attention_heads
      """
      new_shape = tensor.size()[:-1] + (
          num_attention_heads,
          attn_head_size,
      )
      tensor = tensor.view(new_shape)
      if rotary:
          return tensor
      if len(tensor.shape) == 5:
          return tensor.permute(
              0, 1, 3, 2, 4
          )  # (batch, blocks, head, block_length, head_features)
      elif len(tensor.shape) == 4:
          return tensor.permute(
              0, 2, 1, 3
          )  # (batch, head, seq_length, head_features)
      else:
          raise ValueError(
              f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}"
          )
  
imports: 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, CausalLMOutputWithPast, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast
  from ...modeling_utils import PreTrainedModel
  from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
  from ...utils.model_parallel_utils import assert_device_map, get_device_map
  from .configuration_gptj import GPTJConfig
label_desc: Use image node_id 2 for calling the GPTJAttention obj's underlying member method code with example usage: obj._split_heads(tensor, num_attention_heads, attn_head_size, rotary) and returns: tensor, tensor, tensor