Add tanh pre exp

Author: ikostrikov2

flows.py

@@ -64,9 +64,12 @@ def __init__(self,
                  num_hidden,
                  num_cond_inputs=None,
                  s_act='tanh',
-                 t_act='relu'):
+                 t_act='relu',
+                 pre_exp_tanh=False):
         super(MADESplit, self).__init__()
 
+        self.pre_exp_tanh = pre_exp_tanh
+
         activations = {'relu': nn.ReLU, 'sigmoid': nn.Sigmoid, 'tanh': nn.Tanh}
 
         input_mask = get_mask(num_inputs, num_hidden, num_inputs,
@@ -102,6 +105,9 @@ def forward(self, inputs, cond_inputs=None, mode='direct'):
 
             h = self.t_joiner(inputs, cond_inputs)
             a = self.t_trunk(h)
+
+            if self.pre_exp_tanh:
+                a = F.tanh(a)
 
             u = (inputs - m) * torch.exp(-a)
             return u, -a.sum(-1, keepdim=True)
@@ -115,6 +121,9 @@ def forward(self, inputs, cond_inputs=None, mode='direct'):
                 h = self.t_joiner(x, cond_inputs)
                 a = self.t_trunk(h)
 
+                if self.pre_exp_tanh:
+                    a = F.tanh(a)
+
                 x[:, i_col] = inputs[:, i_col] * torch.exp(
                     a[:, i_col]) + m[:, i_col]
             return x, -a.sum(-1, keepdim=True)
@@ -128,7 +137,8 @@ def __init__(self,
                  num_inputs,
                  num_hidden,
                  num_cond_inputs=None,
-                 act='relu'):
+                 act='relu',
+                 pre_exp_tanh=False):
         super(MADE, self).__init__()
 
         activations = {'relu': nn.ReLU, 'sigmoid': nn.Sigmoid, 'tanh': nn.Tanh}