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(
      self, hidden_states: tf.Tensor, rotary: bool
  ) -> tf.Tensor:
      """
      Splits hidden dim into attn_head_size and num_attention_heads
      """
      new_shape = shape_list(hidden_states)[:-1] + [
          self.num_attention_heads,
          self.head_dim,
      ]
      hidden_states = tf.reshape(hidden_states, new_shape)
      if rotary:
          return hidden_states
      if len(shape_list(hidden_states)) == 4:
          return tf.transpose(
              hidden_states, (0, 2, 1, 3)
          )  # (batch, head, seq_length, head_features)
      if len(shape_list(hidden_states)) == 5:
          return tf.transpose(
              hidden_states, (0, 1, 3, 2, 4)
          )  # (batch, blocks, head, block_length, head_features)
      raise ValueError(
          f"Input tensor rank should be one of [4, 5], but is: {len(shape_list(hidden_states))}"
      )
  
imports: from typing import Optional, Tuple, Union
  import numpy
  import tensorflow
  from ...activations_tf import get_tf_activation
  from ...file_utils import DUMMY_INPUTS, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
  from ...modeling_tf_outputs import TFBaseModelOutputWithPast, TFCausalLMOutputWithPast, TFQuestionAnsweringModelOutput, TFSequenceClassifierOutputWithPast
  from ...modeling_tf_utils import TFCausalLanguageModelingLoss, TFModelInputType, TFPreTrainedModel, TFQuestionAnsweringLoss, TFSequenceClassificationLoss, TFSharedEmbeddings, get_initializer, keras_serializable, unpack_inputs
  from ...tf_utils import shape_list, stable_softmax
  from ...utils import logging
  from .configuration_gptj import GPTJConfig
label_desc: Use image node_id 4 for calling the TFGPTJAttention obj's underlying member method code with example usage: obj._split_heads(hidden_states, rotary) and returns: hidden_states, tf, tf