class_docstr: ## Spatial Depth Wise Shared Convolution
  
  We share the same kernel across all channels.
class_docstr_tok: ["#","#","Spatial","Depth","Wise","Shared","Convolution","We","share","the","same","kernel","across","all","channels","."]
docstr: `x` has shape `[seq_len, batch_size, heads, d_k]`
docstr_tok: ["`","x","`","has","shape","`","[","seq_len",",","batch_size",",","heads",",","d_k","]","`"]
code: def forward(self, x: torch.Tensor):
      """
      `x` has shape `[seq_len, batch_size, heads, d_k]`
      """
  
      # Get the shape
      seq_len, batch_size, heads, d_k = x.shape
      # Permute to `[batch_size, heads, d_k, seq_len]`
      x = x.permute(1, 2, 3, 0)
      # Change the shape to `[batch_size * heads * d_k, seq_len]`
      x = x.view(batch_size * heads * d_k, 1, seq_len)
  
      # 1D convolution accepts input of the form `[N, channels, sequence]`
      x = self.conv(x)
      # Crop the right most `kernel_size - 1` results since we padded both sides
      x = x[:, :, : -(self.kernel_size - 1)]
      # Reshape to `[batch_size, heads, d_k, seq_len]`
      x = x.view(batch_size, heads, d_k, seq_len)
      # Permute to `[seq_len, batch_size, heads, d_k]`
      x = x.permute(3, 0, 1, 2)
  
      #
      return x
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/primer_ez/variations.py
imports: import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.transformers import MultiHeadAttention
label_desc: Use image node_id 2 for calling the SpatialDepthWiseSharedConvolution obj's underlying member method code with example usage: obj.forward(x) and returns: x