text: import torch
  from torch import nn
  import torch.nn.functional as F
  from torch.nn.modules.dropout import _DropoutNd
  
  from .swarmlayer import SwarmLayer
  
  from .set_transformer import InducedSetAttentionBlock, RFF
  
  nonlinearities = {}
  nonlinearities['relu']  = nn.ReLU()
  nonlinearities['elu']   = nn.ELU()
  nonlinearities['lrelu'] = nn.LeakyReLU()
  
  
  class MaskedSequential(nn.Module):
      """
      Build a sequential module out of modules that take a mask parameter in their forward() method. Known modules that
      don't take a mask argument (e.g. non-linearities) can be seamlessly included.
      """
      def __init__(self, *mods):
  
          super().__init__()
          self.mods = nn.ModuleList([*mods])
  
      def forward(self, x, mask):
          for m in self.mods:
              print('mod:', type(m))
              if type(m) in [type(nl) for nl in nonlinearities.values()] + \
                      [nn.Linear] + [Dropout2dChannelsLast] + [nn.LSTM]:
                  x = m(x)
              else:
                  x = m(x, mask)
              print('mod:', type(m), 'output:', x)
          return x
  
  
  class SetLinear(nn.Module):
  
      def __init__(self,
                   n_in,
                   n_out,
                   pooling='MEAN'):
          super().__init__()
  
          self.pooling = pooling
          self.ffwd1 = nn.Linear(n_in, n_out)
          self.ffwd2 = nn.Linear(n_in, n_out)
  
      def forward(self, x, mask=None):
          # x is (N, n_samp, n_in)
          N, n_samp, n_in = x.size()
  
          local = self.ffwd1(x)
          glob = self.ffwd2(x)
          if mask is not None:
              mask = mask.unsqueeze(2).float()
              if self.pooling=='MEAN':
                  pool = (glob * mask).sum(dim=1, keepdim=True) / mask.sum(dim=1, keepdim=True)
              else:
                  pool = torch.max(glob+torch.log(mask), dim=1, keepdim=True)[0]
          else:
              if self.pooling=='MEAN':
                  pool = glob.mean(dim=1, keepdim=True)
              else:
                  pool = torch.max(glob, dim=1, keepdim=True)[0]
  
          return local + pool
  
  
  class Dropout2dChannelsLast(_DropoutNd):
  
      def forward(self, input):
          e = torch.ones_like((input[:, :1, :]))
          return input * F.dropout(e, self.p, self.training, self.inplace)
  
  
  def create_model(opt):
  
      nonlinearity = nonlinearities[opt.non_lin]
  
      if opt.type == 'Swarm':
  
          # uses opt. ...
          # n_layers
          # n_in
          # n_out
          # n_hidden
          # n_iter
          # dropout
  
          if opt.n_layers == 1:
              model = SwarmLayer( n_in     = opt.n_in,
                                  n_out    = opt.n_out,
                                  n_hidden = opt.n_hidden,
                                  n_iter   = opt.n_iter,
                                  dropout  = opt.dropout,
                                  n_dim=1, pooling='MEAN', channel_first=False, cache=True)
          else:
              assert opt.n_layers>1
  
              layers = []
              n_out_last = opt.n_in
              for i in range(opt.n_layers-1):
                  layers.append( SwarmLayer(n_in     = n_out_last,
                                            n_out    = opt.n_hidden,
                                            n_hidden = opt.n_hidden,
                                            n_iter   = opt.n_iter,
                                            dropout=opt.dropout,
                                            n_dim=1, pooling='MEAN', channel_first=False, cache=True) )
                  layers.append( nonlinearity)
                  n_out_last = opt.n_hidden
              layers.append( SwarmLayer(n_in     = n_out_last,
                                        n_out    = opt.n_out,
                                        n_hidden = opt.n_hidden,
                                        n_iter   = opt.n_iter,
                                        dropout=opt.dropout,
                                        n_dim=1, pooling='MEAN', channel_first=False, cache=True) )
  
              model = MaskedSequential(*layers)
  
      elif opt.type == 'SetLinear' or opt.type == 'SetLinearMax':
  
          # uses opt. ...
          # n_layers
          # n_in
          # n_out
          # n_hidden
          pooling = 'MEAN' if opt.type == 'SetLinear' else 'MAX'
          layers = []
          n_out_last = opt.n_in
          for i in range(opt.n_layers - 1):
              layers.append(SetLinear(n_in=n_out_last,
                                      n_out=opt.n_hidden,
                                      pooling=pooling) )
              n_out_last = opt.n_hidden
              layers.append(nonlinearity)
          layers.append(SetLinear(n_in=n_out_last,
                                  n_out=opt.n_out,
                                  pooling=pooling) )
  
          model = MaskedSequential(*layers)
  
      elif opt.type == 'SetTransformer':
  
          # uses opt. ...
          # n_layers
          # n_in
          # n_out
          # n_hidden
          # n_heads
          # n_ind_points
  
          d = opt.n_hidden
          h = opt.n_heads
          m = opt.n_ind_points
  
          layers = [nn.Linear( opt.n_in, opt.n_hidden)]
          for _ in range(opt.n_layers):
              layers.append( InducedSetAttentionBlock(d=d, m=m, h=h,first_rff=RFF(d=d),second_rff=RFF(d=d)) )
              if opt.dropout>0.0:
                  layers.append(Dropout2dChannelsLast( p=opt.dropout))
          layers.append(nn.Linear(opt.n_hidden, opt.n_out))
  
          model = MaskedSequential(*layers)
  
      elif opt.type == 'LSTM' or opt.type == 'LSTMS':
  
          # uses opt. ...
          # n_layers
          # n_in
          # n_out
          # n_hidden
  
          # LSTMS is LSTM made set equivariant by sorting
  
          class LSTMModel(nn.Module):
              def __init__(self, opt):
                  super().__init__()
                  self.lstm = nn.LSTM( opt.n_in, opt.n_hidden, opt.n_layers,
                                 batch_first=True, bidirectional=True)
                  self.lin = nn.Linear(opt.n_hidden*2 , opt.n_out)
  
              def forward(self, x, mask):
  
                  if opt.type == 'LSTMS':
                      i = torch.argsort(x[:, :, 0], dim=1)
                      ix = i.unsqueeze(2).expand_as(x)
                      x    = torch.gather(x, dim=1, index=ix)
                      #ask = torch.gather(mask, dim=1, index=i)
  
                  tmp,_ = self.lstm(x)
                  out = self.lin(tmp)
  
                  if opt.type == 'LSTMS':
                      iout = i.unsqueeze(2).expand_as(out)
                      out = torch.zeros_like(out).scatter_(1, iout, out)
  
                  return out
  
          model = LSTMModel(opt)
  
      else:
  
          raise ValueError("Unknown model type {}".format(opt.type))
  
      return model
  