Repository: OpenMined/PySyft
Pull request: cleanup rnn.py
## Description

This pull request simplifies the code in RNNBase. A significant amount of complexity was caused by using the same code to deal handling 2 different data structures of hidden state (LSTMCell has hidden+cell, GRUCell and RNNCell have only hidden). This was simplified by converting these 2 structures into 1 compatible structure internally in the forward method.

I have not added tests. If you think I should add tests for this change, could you point me to an example of how I might do this?

I followed the existing comment structure when making comments.

Fixes #3431 

## Type of change

Please mark the options that are relevant.

- [ ] Added/Modified tutorials
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)

- [x] Reduce code complexity (non-breaking change)

## Checklist:

* [ ] I have added tests for my changes
* [x] I did follow the [contribution guidelines](https://github.com/OpenMined/PySyft/blob/master/CONTRIBUTING.md)
* [ ] I have commented my code following [Google style](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html).


Problem statement:
Reduce complexity of ```forward``` from rnn
Split the code from ```forward``` into two separate functions and remove the ```noqa: C901```.

**Describe alternatives you've considered**
Simplify the function such that no split is required.

**Additional context**
Code quality:
```
2020-04-29T13:13:32.6026063Z ./syft/frameworks/torch/nn/rnn.py:205:5: C901 'RNNBase.forward' is too complex (12)
2020-04-29T13:13:32.6026250Z     def forward(self, x, h=None):
```

Hints from issue discussion:
Hey, is it okay if I start working on this one?
Hey, sure. Assigned the issue to you.

Commit message: cleanup rnn.py

Patch to review:
diff --git a/syft/frameworks/torch/nn/rnn.py b/syft/frameworks/torch/nn/rnn.py
index a5ad317778f..8c79e3b2202 100644
--- a/syft/frameworks/torch/nn/rnn.py
+++ b/syft/frameworks/torch/nn/rnn.py
@@ -185,106 +185,77 @@ def __init__(
         # Link to issue: https://github.com/OpenMined/PySyft/issues/2500
 
         # Build RNN layers
-        self.rnn_forward = nn.ModuleList()
-        for layer in range(self.num_layers):
-            if layer == 0:
-                self.rnn_forward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
-            else:
-                self.rnn_forward.append(base_cell(hidden_size, hidden_size, bias, nonlinearity))
-
+        sizes = [input_size, *[hidden_size] * (self.num_layers - 1)]
+        self.rnn_forward = nn.ModuleList(
+            (base_cell(sz, hidden_size, bias, nonlinearity) for sz in sizes)
+        )
         if self.bidirectional:
-            self.rnn_backward = nn.ModuleList()
-            for layer in range(self.num_layers):
-                if layer == 0:
-                    self.rnn_backward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
-                else:
-                    self.rnn_backward.append(
-                        base_cell(hidden_size, hidden_size, bias, nonlinearity)
-                    )
+            self.rnn_backward = nn.ModuleList(
+                (base_cell(sz, hidden_size, bias, nonlinearity) for sz in sizes)
+            )
 
-    def forward(self, x, h=None):
-        # If batch_first == True, swap axis with seq_len
+    def forward(self, x, hc=None):
+        # Treat hc[0] as the hidden state regardless of whether it is an LSTM or GRU
+        if not self.is_lstm and hc is not None:
+            hc = [hc]
+
+        # If batch_first == True, swap batch with seq_len
         # At the end of the process we swap it back to the original structure
         if self.batch_first:
-            x, h = self._swap_axis(x, h)
-
-        # If it is LSTM, get hidden and cell states
-        if h is not None and self.is_lstm:
-            h, c = h
+            x = torch.transpose(x, 0, 1)
+            if hc is not None:
+                hc = [torch.transpose(item, 0, 1) for item in hc]
 
         batch_size = x.shape[1]
         seq_len = x.shape[0]
 
         # Initiate states if needed
-        if h is None:
-            h = self._init_hidden(x)
-            c = self._init_hidden(x) if self.is_lstm else None
+        if hc is None:
+            hc = [self._init_hidden(x) for i in range(2 if self.is_lstm else 1)]
 
-        # If bidirectional==True, split states in two, each one for each direction
+        # If bidirectional==True, split states in two, one for each direction
         if self.bidirectional:
-            h = h.contiguous().view(self.num_layers, 2, batch_size, self.hidden_size)
-            h_for = h[:, 0, :, :]
-            h_back = h[:, 1, :, :]
-            if self.is_lstm:
-                c = c.contiguous().view(self.num_layers, 2, batch_size, self.hidden_size)
-                c_for = c[:, 0, :, :]
-                c_back = c[:, 1, :, :]
-            else:
-                c_for = c_back = None
-
+            hc = [
+                item.contiguous().view(self.num_layers, 2, batch_size, self.hidden_size)
+                for item in hc
+            ]
+            hc_fwd = [item[:, 0, :, :] for item in hc]
+            hc_back = [item[:, 1, :, :] for item in hc]
         else:
-            h_for = h
-            c_for = c if self.is_lstm else None
+            hc_fwd = hc
 
         # Run through rnn in the forward direction
         output = x.new(seq_len, batch_size, self.hidden_size).zero_()
         for t in range(seq_len):
-            h_for, c_for = self._apply_time_step(x, h_for, c_for, t)
-            output[t, :, :] = h_for[-1, :, :]
-
-        hidden = h_for
-        cell = c_for  # None if it is not an LSTM
+            hc_fwd = self._apply_time_step(x, hc_fwd, t)
+            output[t, :, :] = hc_fwd[0][-1, :, :]
 
         # Run through rnn in the backward direction if bidirectional==True
         if self.bidirectional:
             output_back = x.new(seq_len, batch_size, self.hidden_size).zero_()
             for t in range(seq_len - 1, -1, -1):
-                h_back, c_back = self._apply_time_step(x, h_back, c_back, t, reverse_direction=True)
-                output_back[t, :, :] = h_back[-1, :, :]
+                hc_back = self._apply_time_step(x, hc_back, t, reverse_direction=True)
+                output_back[t, :, :] = hc_back[0][-1, :, :]
 
             # Concatenate both directions
             output = torch.cat((output, output_back), dim=-1)
-            hidden = torch.cat((hidden, h_back), dim=0)
-            if self.is_lstm:
-                cell = torch.cat((cell, c_back), dim=0)
+            hidden = [
+                torch.cat((hid_item, back_item), dim=0)
+                for hid_item, back_item in zip(hc_fwd, hc_back)
+            ]
+        else:
+            hidden = hc_fwd
 
         # If batch_first == True, swap axis back to get original structure
         if self.batch_first:
             output = torch.transpose(output, 0, 1)
-            hidden = torch.transpose(hidden, 0, 1)
-            if self.is_lstm:
-                cell = torch.transpose(cell, 0, 1)
+            hidden = [torch.transpose(item, 0, 1) for item in hidden]
 
-        hidden = (hidden, cell) if self.is_lstm else hidden
+        # Reshape hidden to the original shape of h
+        hidden = tuple(hidden) if self.is_lstm else hidden[0]
 
         return output, hidden
 
-    def _swap_axis(self, x, h):
-        """
-        This method swap the axes for batch_size and seq_len. It is used when
-        batch_first==True.
-        """
-        x = torch.transpose(x, 0, 1)
-        if h is not None:
-            if self.is_lstm:
-                h, c = h
-                h = torch.transpose(h, 0, 1)
-                c = torch.transpose(c, 0, 1)
-                h = (h, c)
-            else:
-                h = torch.transpose(h, 0, 1)
-        return x, h
-
     def _init_hidden(self, input):
         """
         This method initializes a hidden state when no hidden state is provided
@@ -309,34 +280,24 @@ def _init_hidden(self, input):
                 h = h.share(*owners, crypto_provider=crypto_provider)
         return h
 
-    def _apply_time_step(self, x, h, c, t, reverse_direction=False):
+    def _apply_time_step(self, x, h, t, reverse_direction=False):
         """
         Apply RNN layers at time t, given input and previous hidden states
         """
         rnn_layers = self.rnn_backward if reverse_direction else self.rnn_forward
 
-        h_next = h.new(h.shape).zero_()
-        c_next = c.new(c.shape).zero_() if self.is_lstm else None
+        h = torch.stack([*h])
+        h_next = torch.stack([item.new(item.shape).zero_() for item in h])
 
         for layer in range(self.num_layers):
-            if layer == 0:
-                if self.is_lstm:
-                    h_next[layer, :, :], c_next[layer, :, :] = rnn_layers[layer](
-                        x[t, :, :], (h[layer, :, :], c[layer, :, :])
-                    )
-                else:
-                    h_next[layer, :, :] = rnn_layers[layer](x[t, :, :], h[layer, :, :])
+            inp = x[t, :, :] if layer == 0 else h_next[0][layer - 1, :, :].clone()
+
+            if self.is_lstm:
+                h_next[:, layer, :, :] = torch.stack(rnn_layers[layer](inp, h[:, layer, :, :]))
             else:
-                if self.is_lstm:
-                    h_next[layer, :, :], c_next[layer, :, :] = rnn_layers[layer](
-                        h_next[layer - 1, :, :].clone(), (h[layer, :, :], c[layer, :, :])
-                    )
-                else:
-                    h_next[layer, :, :] = rnn_layers[layer](
-                        h_next[layer - 1, :, :].clone(), h[layer, :, :]
-                    )
-
-        return h_next, c_next
+                h_next[0][layer, :, :] = rnn_layers[layer](inp, h[0][layer, :, :])
+
+        return h_next
 
 
 class RNN(RNNBase):


Code review comments:
diff_hunk: @@ -185,106 +185,77 @@ def __init__(
         # Link to issue: https://github.com/OpenMined/PySyft/issues/2500
 
         # Build RNN layers
-        self.rnn_forward = nn.ModuleList()
-        for layer in range(self.num_layers):
-            if layer == 0:
-                self.rnn_forward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
-            else:
-                self.rnn_forward.append(base_cell(hidden_size, hidden_size, bias, nonlinearity))
-
+        sizes = [input_size, *[hidden_size] * (self.num_layers - 1)]
+        self.rnn_forward = nn.ModuleList(
+            (base_cell(sz, hidden_size, bias, nonlinearity) for sz in sizes)
+        )
         if self.bidirectional:
-            self.rnn_backward = nn.ModuleList()
-            for layer in range(self.num_layers):
-                if layer == 0:
-                    self.rnn_backward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
-                else:
-                    self.rnn_backward.append(
-                        base_cell(hidden_size, hidden_size, bias, nonlinearity)
-                    )
+            self.rnn_backward = nn.ModuleList(
+                (base_cell(sz, hidden_size, bias, nonlinearity) for sz in sizes)
+            )
 
-    def forward(self, x, h=None):
-        # If batch_first == True, swap axis with seq_len
+    def forward(self, x, hc=None):
+        # Treat hc[0] as the hidden state regardless of whether it is an LSTM or GRU
+        if not self.is_lstm and hc is not None:
+            hc = [hc]
+
+        # If batch_first == True, swap batch with seq_len
         # At the end of the process we swap it back to the original structure
         if self.batch_first:
-            x, h = self._swap_axis(x, h)
-
-        # If it is LSTM, get hidden and cell states
-        if h is not None and self.is_lstm:
-            h, c = h
+            x = torch.transpose(x, 0, 1)
+            if hc is not None:
+                hc = [torch.transpose(item, 0, 1) for item in hc]
 
         batch_size = x.shape[1]
         seq_len = x.shape[0]
 
         # Initiate states if needed
-        if h is None:
-            h = self._init_hidden(x)
-            c = self._init_hidden(x) if self.is_lstm else None
+        if hc is None:
+            hc = [self._init_hidden(x) for i in range(2 if self.is_lstm else 1)]
path: syft/frameworks/torch/nn/rnn.py
text: @user1:
i is not used, you can write for _ in range(...) to make it clear

@author:
Oops, yep.

diff_hunk: @@ -185,106 +185,77 @@ def __init__(
         # Link to issue: https://github.com/OpenMined/PySyft/issues/2500
 
         # Build RNN layers
-        self.rnn_forward = nn.ModuleList()
-        for layer in range(self.num_layers):
-            if layer == 0:
-                self.rnn_forward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
-            else:
-                self.rnn_forward.append(base_cell(hidden_size, hidden_size, bias, nonlinearity))
-
+        sizes = [input_size, *[hidden_size] * (self.num_layers - 1)]
+        self.rnn_forward = nn.ModuleList(
+            (base_cell(sz, hidden_size, bias, nonlinearity) for sz in sizes)
+        )
         if self.bidirectional:
-            self.rnn_backward = nn.ModuleList()
-            for layer in range(self.num_layers):
-                if layer == 0:
-                    self.rnn_backward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
-                else:
-                    self.rnn_backward.append(
-                        base_cell(hidden_size, hidden_size, bias, nonlinearity)
-                    )
+            self.rnn_backward = nn.ModuleList(
+                (base_cell(sz, hidden_size, bias, nonlinearity) for sz in sizes)
+            )
 
-    def forward(self, x, h=None):
-        # If batch_first == True, swap axis with seq_len
+    def forward(self, x, hc=None):
+        # Treat hc[0] as the hidden state regardless of whether it is an LSTM or GRU
+        if not self.is_lstm and hc is not None:
+            hc = [hc]
+
+        # If batch_first == True, swap batch with seq_len
         # At the end of the process we swap it back to the original structure
         if self.batch_first:
-            x, h = self._swap_axis(x, h)
-
-        # If it is LSTM, get hidden and cell states
-        if h is not None and self.is_lstm:
-            h, c = h
+            x = torch.transpose(x, 0, 1)
+            if hc is not None:
path: syft/frameworks/torch/nn/rnn.py
text: @user1:
I think it would be clearer if hc is processed all at once (in the same place as hc = [hc]).
For instance, we could have
```
if hc is not None:
    if not self.is_lstm:
        hc = [hc]
    if self.batch_first:
       hc = [torch.transpose(item, 0, 1) for item in hc]
```
at the top of the function

@author:
I agree, and also added comments here to clarify what `hc` represents

diff_hunk: @@ -185,106 +185,77 @@ def __init__(
         # Link to issue: https://github.com/OpenMined/PySyft/issues/2500
 
         # Build RNN layers
-        self.rnn_forward = nn.ModuleList()
-        for layer in range(self.num_layers):
-            if layer == 0:
-                self.rnn_forward.append(base_cell(input_size, hidden_size, bias, nonlinearity))
-            else:
-                self.rnn_forward.append(base_cell(hidden_size, hidden_size, bias, nonlinearity))
-
+        sizes = [input_size, *[hidden_size] * (self.num_layers - 1)]
path: syft/frameworks/torch/nn/rnn.py
text: @user1:
I think `sizes = [input_size, *(hidden_size for _ in range(self.num_layers - 1))]` is slightly clearer but maybe it's just my opinion

@author:
Agreed

diff_hunk: @@ -309,34 +280,24 @@ def _init_hidden(self, input):
                 h = h.share(*owners, crypto_provider=crypto_provider)
         return h
 
-    def _apply_time_step(self, x, h, c, t, reverse_direction=False):
+    def _apply_time_step(self, x, h, t, reverse_direction=False):
         """
         Apply RNN layers at time t, given input and previous hidden states
         """
         rnn_layers = self.rnn_backward if reverse_direction else self.rnn_forward
 
-        h_next = h.new(h.shape).zero_()
-        c_next = c.new(c.shape).zero_() if self.is_lstm else None
+        h = torch.stack([*h])
+        h_next = torch.stack([item.new(item.shape).zero_() for item in h])
path: syft/frameworks/torch/nn/rnn.py
text: @user1:
I'm not sure what happens here but I don't think it's the easiest way to do it... Is it?

@author:
(recall `h` is now `hc`, per previous comment)
`torch.stack([*hc])` just ensures that `hc` is a tensor, not a list.
And no, the second line is convoluted per my blindly direct translation of the previous code. I've now replaced it with: `hc_next = torch.zeros_like(hc)`

diff_hunk: @@ -309,34 +280,24 @@ def _init_hidden(self, input):
                 h = h.share(*owners, crypto_provider=crypto_provider)
         return h
 
-    def _apply_time_step(self, x, h, c, t, reverse_direction=False):
+    def _apply_time_step(self, x, h, t, reverse_direction=False):
path: syft/frameworks/torch/nn/rnn.py
text: @user1:
Can you call the argument `hc` to be coherent with what you wrote above?

@author:
Oops, yep.