code: def split_heads(self, x, k=False):
      new_x_shape = x.size()[:-1] + (
          self.n_head,
          x.size(-1) // self.n_head,
      )
      x = x.view(*new_x_shape)  # in Tensorflow implem: fct split_states
      if k:
          return x.permute(0, 2, 3, 1)
      else:
          return x.permute(0, 2, 1, 3)
  
imports: from __future__ import with_statement
  from __future__ import division
  from __future__ import absolute_import
  import copy
  import io
  import json
  import logging
  import math
  import re
  import tarfile
  import numpy
  import torch
  from torch.nn import Parameter
  from allennlp.common.checks import ConfigurationError
  from allennlp.common.file_utils import cached_path
  from allennlp.common.from_params import FromParams
  from io import open
label_desc: Use image node_id 4 for calling the Attention obj's underlying member method code with example usage: obj.split_heads(x, k) and returns: x, x