text: import torch.nn as nn
  import torch
  import torch.nn.init as init
  
  class Linear(nn.Module):
      ''' Simple Linear layer with xavier init '''
      def __init__(self, d_in, d_out, bias=True):
          super(Linear, self).__init__()
          self.linear = nn.Linear(d_in, d_out, bias=bias)
          init.xavier_normal(self.linear.weight)
  
      def forward(self, x):
          return self.linear(x)
  
  class Bottle(nn.Module):
      ''' Perform the reshape routine before and after an operation '''
  
      def forward(self, input):
          if len(input.size()) <= 2:
              return super(Bottle, self).forward(input)
          size = input.size()[:2]
          out = super(Bottle, self).forward(input.view(size[0]*size[1], -1))
          return out.view(size[0], size[1], -1)
  
  class BottleLinear(Bottle, Linear):
      ''' Perform the reshape routine before and after a linear projection '''
      pass
  