code: def _merge_heads(
      x: torch.Tensor, num_heads: int, head_dim: int
  ) -> torch.Tensor:
      """
      Merge heads together over the last dimenstion
  
      Args:
          x: (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim]
  
      Returns:
          torch.tensor: [batch_size, seq_length, num_heads * head_dim]
      """
      # What we want to achieve is:
      # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim
      batch_size_and_num_heads, seq_length, _ = x.shape
      batch_size = batch_size_and_num_heads // num_heads
  
      # First view to decompose the batch size
      # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim
      x = x.view(batch_size, num_heads, seq_length, head_dim)
  
      # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim
      x = x.permute(0, 2, 1, 3)
  
      # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim
      return x.reshape(batch_size, seq_length, num_heads * head_dim)
  
imports: import math
  import os
  import warnings
  from typing import Optional, Tuple, Union
  import torch
  import torch.distributed
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import LayerNorm
  from torch.nn import functional
  from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions
  from transformers import BloomConfig, PreTrainedModel
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelRowLinear, TensorParallelHead
label_desc: Use image node_id 6 for calling a global function with example usage: _merge_heads(x, num_heads, head_dim) and returns: x