class_docstr: ## Spatial Depth Wise Convolution
class_docstr_tok: ["#","#","Spatial","Depth","Wise","Convolution"]
docstr: * `d_k` is the number of channels in each head
docstr_tok: ["*","`","d_k","`","is","the","number","of","channels","in","each","head"]
code: def __init__(self, d_k: int, kernel_size: int = 3):
      """
      * `d_k` is the number of channels in each head
      """
      super().__init__()
      self.kernel_size = kernel_size
      # We use PyTorch's `Conv1d` module.
      # We set the number of groups to be equal to the number of channels so that it does a separate convolution
      # (with different kernels) for each channel.
      # We add padding to both sides and later crop the right most `kernel_size - 1` results
      self.conv = nn.Conv1d(
          in_channels=d_k,
          out_channels=d_k,
          kernel_size=(kernel_size,),
          padding=(kernel_size - 1,),
          groups=d_k,
      )
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/primer_ez/__init__.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 1 to create a new SpatialDepthWiseConvolution object from inherited base classes: Module with example: obj = SpatialDepthWiseConvolution(d_k, kernel_size)