def METHOD_NAME(self, x):
    "return shape [`batch`, `head`, `sequence`, `features`]"
    new_shape = x.size()[:-1] + (self.n_head, x.size(-1)//self.n_head)
    x = x.view(*new_shape)
    return x.permute(0, 2, 1, 3)