text: """
  indexed.py
  ========
  Contain the core functions for the indexed operations
  """
  
  import logging
  import math
  
  import torch
  import torch.nn as nn
  from torch.nn.parameter import Parameter
  
  import indexedconv.utils as utils
  
  
  class IndexedMaxPool2d(nn.Module):
      r"""
      Compute the Max Pooling 2d operation on a batch of features of vector images wrt a matrix of indices
  
      Args:
          indices (LongTensor): index tensor of shape (L x kernel_size),
              having on each row the indices of neighbors of each element of the input
              a -1 indicates the absence of a neighbor, which is handled as zero-padding
      Shape:
          - Input: :math:`(N, C, L_{in})`
          - Output: :math:`(N, C, L_{out})`
      """
      def __init__(self, indices):
          super(IndexedMaxPool2d, self).__init__()
          self.logger = logging.getLogger(__name__ + '.IndexedMaxPool2d')
          self.indices = indices
          self.indices, self.mask = utils.prepare_mask(self.indices)
          self.register_buffer('indices_', self.indices)
          self.register_buffer('mask_', self.mask)
  
      def forward(self, input_images):
          self.logger.debug('Max pool image')
  
          col = input_images[..., self.indices_] * self.mask_
  
          out, _ = torch.max(col, 2)
  
          return out
  
  
  class IndexedAveragePool2d(nn.Module):
      r"""
      Compute the Average Pooling 2d operation on a batch of features of vector images wrt a matrix of indices
  
      Args:
          indices (LongTensor): index tensor of shape (L x kernel_size), having on each
                                row the indices of neighbors of each element of the input
                                a -1 indicates the absence of a neighbor, which is handled
                                as zero-padding
      Shape:
          - Input: :math:`(N, C, L_{in})`
          - Output: :math:`(N, C, L_{out})`
  
      """
      def __init__(self, indices):
          super(IndexedAveragePool2d, self).__init__()
          self.logger = logging.getLogger(__name__ + '.IndexedAveragePool2d')
          self.indices = indices
          self.indices, self.mask = utils.prepare_mask(self.indices)
          self.register_buffer('indices_', self.indices)
          self.register_buffer('mask_', self.mask)
  
      def forward(self, input_images):
          self.logger.debug('Average pool image')
  
          col = input_images[..., self.indices_] * self.mask_
  
          out = torch.mean(col, 2)
  
          return out
  
  
  class IndexedConv(nn.Module):
      r"""Applies a convolution over an input tensor where neighborhood relationships
      between elements are explicitly provided via an `indices` tensor.
  
      The output value of the layer with input size :math:`(N, C_{in}, L)` and output
      :math:`(N, C_{out}, L)` can be described as:
  
      .. math::
  
          \begin{array}{ll}
          out(N_i, C_{out_j})  = bias(C_{out_j})
                         + \sum_{{c}=0}^{C_{in}-1}
                           \sum_{{i}=0}^{L-1}
                           \sum_{{k}=0}^{K} weight(C_{out_j}, c, k) * input(N_i, c, indices(i, k))
          \end{array}
  
      where
  
      | `indices` is a L x K tensor, where `K` is the size of the convolution kernel,
      | providing the indices of the `K` neighbors of input element `i`.
      | A -1 entry means zero-padding.
  
      Args:
          in_channels (int): Number of channels in the input tensor
  
          out_channels (int): Number of channels produced by the convolution
  
          indices (LongTensor): index tensor of shape (L x kernel_size), having on each
          row the indices of neighbors of each element of the input a -1 indicates the absence of a
          neighbor, which is handled as zero-padding
  
          bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``True``
  
      Shape:
          - Input: :math:`(N, C_{in}, L)`
          - Output: :math:`(N, C_{out}, L)`
  
      Attributes:
          weight (Tensor): the learnable weights of the module of shape
              (out_channels, in_channels, kernel_size)
          bias (Tensor):   the learnable bias of the module of shape
              (out_channels)
  
      Examples::
  
          >>> indices = (10 * torch.rand(50, 3)).type(torch.LongTensor)
          >>> m = nn.IndexedConv(16, 3, indices)
          >>> input = torch.randn(20, 16, 50)
          >>> output = m(input)
      """
  
      def __init__(self, in_channels, out_channels, indices, bias=True):
          super(IndexedConv, self).__init__()
          self.logger = logging.getLogger(__name__ + '.IndexedConv')
  
          groups = 1
  
          kernel_size = indices.shape[0]
  
          self.in_channels = in_channels
          self.out_channels = out_channels
          self.kernel_size = kernel_size
          self.indices, self.mask = utils.prepare_mask(indices)
          self.output_width = self.indices.shape[1]
  
          self.register_buffer('indices_', self.indices)
          self.register_buffer('mask_', self.mask)
  
          self.weight = Parameter(torch.Tensor(
              out_channels, in_channels // groups, kernel_size))
  
          if bias:
              self.bias = Parameter(torch.Tensor(out_channels))
          else:
              self.register_parameter('bias', None)
  
          self.reset_parameters()
  
      def reset_parameters(self):
          n = self.in_channels * self.kernel_size
          stdv = 1. / math.sqrt(n)
          self.weight.data.uniform_(-stdv, stdv)
          if self.bias is not None:
              self.bias.data.uniform_(-stdv, stdv)
  
      def forward(self, input):
          nbatch = input.shape[0]
          col = input[..., self.indices_] * self.mask_
          # col is of shape (N, C_in, K, Wo)
          col = col.reshape(nbatch, -1, self.output_width)
          out = torch.bmm(self.weight.view(self.out_channels, -1).expand(nbatch, -1, -1), col)
          if self.bias is not None:
              out = out + self.bias.unsqueeze(1)
  
          return out
  
      def __repr__(self):
          s = '{name} ({in_channels}, {out_channels}, kernel_size={kernel_size}'
          if self.bias is None:
              s += ', bias=False'
          s += ')'
          return s.format(name=self.__class__.__name__, **self.__dict__)
  