text: import functools
  
  import torch
  from torch import nn
  import torch.nn.functional as F
  from torch.nn.modules.conv import _ConvNd
  from torch.nn.modules.utils import _pair, _triple
  from torch.nn.parameter import Parameter
  
  
  class _routing(nn.Module):
  
      def __init__(self, in_channels, num_experts, dropout_rate):
          super(_routing, self).__init__()
          
          self.dropout = nn.Dropout(dropout_rate)
          self.fc = nn.Linear(in_channels, num_experts)
  
      def forward(self, x, embedding=None):
          bs, cs = x.shape[0], x.shape[1]
          x = x.reshape(bs, cs)
          x = self.dropout(x)
          if embedding is None:
              x = self.fc(x)
          else:
              x = torch.cat([x, embedding], axis=1)
              x = self.fc(x)
          return F.sigmoid(x)
      
  
  class CondConv3D(_ConvNd):
      r"""Learn specialized convolutional kernels for each example.
      As described in the paper
      `CondConv: Conditionally Parameterized Convolutions for Efficient Inference`_ ,
      conditionally parameterized convolutions (CondConv), 
      which challenge the paradigm of static convolutional kernels 
      by computing convolutional kernels as a function of the input.
      Args:
          in_channels (int): Number of channels in the input image
          out_channels (int): Number of channels produced by the convolution
          kernel_size (int or tuple): Size of the convolving kernel
          stride (int or tuple, optional): Stride of the convolution. Default: 1
          padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
          padding_mode (string, optional): ``'zeros'``, ``'reflect'``, ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
          dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
          groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
          bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
          num_experts (int): Number of experts per layer 
          embeddings (int, optional): If >0, then feed in embedding feature into the routing_fn. Default: 0 
      Shape:
          - Input: :math:`(N, C_{in}, H_{in}, W_{in})`
          - Output: :math:`(N, C_{out}, H_{out}, W_{out})` where
            .. math::
                H_{out} = \left\lfloor\frac{H_{in}  + 2 \times \text{padding}[0] - \text{dilation}[0]
                          \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
            .. math::
                W_{out} = \left\lfloor\frac{W_{in}  + 2 \times \text{padding}[1] - \text{dilation}[1]
                          \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
      Attributes:
          weight (Tensor): the learnable weights of the module of shape
                           :math:`(\text{out\_channels}, \frac{\text{in\_channels}}{\text{groups}},`
                           :math:`\text{kernel\_size[0]}, \text{kernel\_size[1]})`.
                           The values of these weights are sampled from
                           :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
                           :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
          bias (Tensor):   the learnable bias of the module of shape (out_channels). If :attr:`bias` is ``True``,
                           then the values of these weights are
                           sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
                           :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
      .. _CondConv: Conditionally Parameterized Convolutions for Efficient Inference:
         https://arxiv.org/abs/1904.04971
      """
  
      def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                   padding=0, dilation=1, groups=1, embeddings = 0,
                   bias=True, padding_mode='zeros', num_experts=3, dropout_rate=0.2):
          kernel_size = _triple(kernel_size)
          stride = _triple(stride)
          padding = _triple(padding)
          dilation = _triple(dilation)
          super(CondConv3D, self).__init__(
              in_channels, out_channels, kernel_size, stride, padding, dilation,
              False, _pair(0), groups, bias, padding_mode)
  
          self._avg_pooling = functools.partial(F.adaptive_avg_pool3d, output_size=(1, 1, 1))
          routing_channels = in_channels + embeddings
          self._routing_fn = _routing(routing_channels, num_experts, dropout_rate)
          
          self.weight = Parameter(torch.Tensor(
              num_experts, out_channels, in_channels // groups, *kernel_size))
          self.init_weights()
          # self.reset_parameters()
  
  
      def init_weights(self):
          nn.init.xavier_normal_(self.weight)
          nn.init.constant_(self.bias, 0)
          
      def _conv_forward(self, input, weight):
          if self.padding_mode != 'zeros':
              return F.conv3d(F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
                              weight, self.bias, self.stride,
                              _pair(0), self.dilation, self.groups)
          return F.conv3d(input, weight, self.bias, self.stride,
                          self.padding, self.dilation, self.groups)
      
      def forward(self, inputs, embeddings=None):
          b, _, _, _, _ = inputs.size()
          res = []
          idx = 0
          pooled_inputs = self._avg_pooling(inputs)
          routing_weights = self._routing_fn(pooled_inputs, embeddings)
          # print((routing_weights[: , :, None, None, None, None, None] * self.weight).shape)
          kernels = torch.sum(routing_weights[: , :, None, None, None, None, None] * self.weight, 1)
          for i in range(b):
              out = self._conv_forward(inputs[i].unsqueeze(0), kernels[i])
              res.append(out)
          idx += 1
          return torch.cat(res, dim=0)
  
  
  
  
  
  # class _routing(nn.Module):
  
  #     def __init__(self, in_channels, num_experts, dropout_rate):
  #         super(_routing, self).__init__()
          
  #         self.dropout = nn.Dropout(dropout_rate)
  #         self.fc = nn.Linear(in_channels, num_experts)
  
  #     def forward(self, x, embedding=None):
  #         x = torch.flatten(x)
  #         x = self.dropout(x)
  #         if embedding is None:
  #             x = self.fc(x)
  #         else:
  #             x = torch.cat([x, embedding])
  #             x = self.fc(x)
  #         return F.sigmoid(x)
      
  
  # class CondConv3D(_ConvNd):
  #     r"""Learn specialized convolutional kernels for each example.
  #     As described in the paper
  #     `CondConv: Conditionally Parameterized Convolutions for Efficient Inference`_ ,
  #     conditionally parameterized convolutions (CondConv), 
  #     which challenge the paradigm of static convolutional kernels 
  #     by computing convolutional kernels as a function of the input.
  #     Args:
  #         in_channels (int): Number of channels in the input image
  #         out_channels (int): Number of channels produced by the convolution
  #         kernel_size (int or tuple): Size of the convolving kernel
  #         stride (int or tuple, optional): Stride of the convolution. Default: 1
  #         padding (int or tuple, optional): Zero-padding added to both sides of the input. Default: 0
  #         padding_mode (string, optional): ``'zeros'``, ``'reflect'``, ``'replicate'`` or ``'circular'``. Default: ``'zeros'``
  #         dilation (int or tuple, optional): Spacing between kernel elements. Default: 1
  #         groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
  #         bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
  #         num_experts (int): Number of experts per layer 
  #         embeddings (int, optional): If >0, then feed in embedding feature into the routing_fn. Default: 0 
  #     Shape:
  #         - Input: :math:`(N, C_{in}, H_{in}, W_{in})`
  #         - Output: :math:`(N, C_{out}, H_{out}, W_{out})` where
  #           .. math::
  #               H_{out} = \left\lfloor\frac{H_{in}  + 2 \times \text{padding}[0] - \text{dilation}[0]
  #                         \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
  #           .. math::
  #               W_{out} = \left\lfloor\frac{W_{in}  + 2 \times \text{padding}[1] - \text{dilation}[1]
  #                         \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
  #     Attributes:
  #         weight (Tensor): the learnable weights of the module of shape
  #                          :math:`(\text{out\_channels}, \frac{\text{in\_channels}}{\text{groups}},`
  #                          :math:`\text{kernel\_size[0]}, \text{kernel\_size[1]})`.
  #                          The values of these weights are sampled from
  #                          :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
  #                          :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
  #         bias (Tensor):   the learnable bias of the module of shape (out_channels). If :attr:`bias` is ``True``,
  #                          then the values of these weights are
  #                          sampled from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
  #                          :math:`k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}`
  #     .. _CondConv: Conditionally Parameterized Convolutions for Efficient Inference:
  #        https://arxiv.org/abs/1904.04971
  #     """
  
  #     def __init__(self, in_channels, out_channels, kernel_size, stride=1,
  #                  padding=0, dilation=1, groups=1, embeddings = 0,
  #                  bias=True, padding_mode='zeros', num_experts=3, dropout_rate=0.2):
  #         kernel_size = _triple(kernel_size)
  #         stride = _triple(stride)
  #         padding = _triple(padding)
  #         dilation = _triple(dilation)
  #         super(CondConv3D, self).__init__(
  #             in_channels, out_channels, kernel_size, stride, padding, dilation,
  #             False, _pair(0), groups, bias, padding_mode)
  
  #         self._avg_pooling = functools.partial(F.adaptive_avg_pool3d, output_size=(1, 1, 1))
  #         routing_channels = in_channels + embeddings
  #         self._routing_fn = _routing(routing_channels, num_experts, dropout_rate)
          
  #         self.weight = Parameter(torch.Tensor(
  #             num_experts, out_channels, in_channels // groups, *kernel_size))
  #         self.init_weights()
  #         # self.reset_parameters()
  
  
  #     def init_weights(self):
  #         nn.init.xavier_normal_(self.weight)
  #         nn.init.constant_(self.bias, 0)
          
  #     def _conv_forward(self, input, weight):
  #         if self.padding_mode != 'zeros':
  #             return F.conv3d(F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
  #                             weight, self.bias, self.stride,
  #                             _pair(0), self.dilation, self.groups)
  #         return F.conv3d(input, weight, self.bias, self.stride,
  #                         self.padding, self.dilation, self.groups)
      
  #     def forward(self, inputs, embeddings=None):
  #         b, _, _, _, _ = inputs.size()
  #         res = []
  #         idx = 0
  #         for input in inputs:
  #             input = input.unsqueeze(0)
  #             pooled_inputs = self._avg_pooling(input)
  #             routing_weights = self._routing_fn(pooled_inputs, embeddings[idx])
  #             kernels = torch.sum(routing_weights[: ,None, None, None, None, None] * self.weight, 0)
  #             out = self._conv_forward(input, kernels)
  #             res.append(out)
  #             idx += 1
  #         return torch.cat(res, dim=0)