code: def split_states(x, n):
      """Reshape the last dimension of x into [n, x.shape[-1]/n]."""
      *start, m = shape_list(x)
      return tf.reshape(x, start + [n, m // n])
  
file_path: finetune/finetune/base_models/gpt2/featurizer.py
imports: import functools
  import tensorflow
  from finetune.util.shapes import shape_list, lengths_from_eos_idx
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.nn.activations import gelu
  from finetune.base_models.gpt.featurizer import norm, dropout, get_pos_values
label_desc: Use image node_id 2 for calling a global function with example usage: split_states(x, n) and returns: tf