We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Is StackedLSTMCell correct implementation?
For enas, StackedLSTMCell is implemented as below
 https://github.com/microsoft/nni/blob/bf8be1e705cdc0c153d6f38d1a83bd3fe55d7a1e/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py#L23-L28

nn.LSTMCell outputs (hidden state, cell state) but the order of output is changed in the implementation.
As a result, the cell state is used for output of LSTM.

Is this an intention or wrong implementation?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/sdk/pynni/nni/nas/pytorch/enas/mutator.py`
Content:
```
1 # Copyright (c) Microsoft Corporation.
2 # Licensed under the MIT license.
3 
4 import torch
5 import torch.nn as nn
6 import torch.nn.functional as F
7 
8 from nni.nas.pytorch.mutator import Mutator
9 from nni.nas.pytorch.mutables import LayerChoice, InputChoice, MutableScope
10 
11 
12 class StackedLSTMCell(nn.Module):
13     def __init__(self, layers, size, bias):
14         super().__init__()
15         self.lstm_num_layers = layers
16         self.lstm_modules = nn.ModuleList([nn.LSTMCell(size, size, bias=bias)
17                                            for _ in range(self.lstm_num_layers)])
18 
19     def forward(self, inputs, hidden):
20         prev_c, prev_h = hidden
21         next_c, next_h = [], []
22         for i, m in enumerate(self.lstm_modules):
23             curr_c, curr_h = m(inputs, (prev_c[i], prev_h[i]))
24             next_c.append(curr_c)
25             next_h.append(curr_h)
26             # current implementation only supports batch size equals 1,
27             # but the algorithm does not necessarily have this limitation
28             inputs = curr_h[-1].view(1, -1)
29         return next_c, next_h
30 
31 
32 class EnasMutator(Mutator):
33     """
34     A mutator that mutates the graph with RL.
35 
36     Parameters
37     ----------
38     model : nn.Module
39         PyTorch model.
40     lstm_size : int
41         Controller LSTM hidden units.
42     lstm_num_layers : int
43         Number of layers for stacked LSTM.
44     tanh_constant : float
45         Logits will be equal to ``tanh_constant * tanh(logits)``. Don't use ``tanh`` if this value is ``None``.
46     cell_exit_extra_step : bool
47         If true, RL controller will perform an extra step at the exit of each MutableScope, dump the hidden state
48         and mark it as the hidden state of this MutableScope. This is to align with the original implementation of paper.
49     skip_target : float
50         Target probability that skipconnect will appear.
51     temperature : float
52         Temperature constant that divides the logits.
53     branch_bias : float
54         Manual bias applied to make some operations more likely to be chosen.
55         Currently this is implemented with a hardcoded match rule that aligns with original repo.
56         If a mutable has a ``reduce`` in its key, all its op choices
57         that contains `conv` in their typename will receive a bias of ``+self.branch_bias`` initially; while others
58         receive a bias of ``-self.branch_bias``.
59     entropy_reduction : str
60         Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced.
61     """
62 
63     def __init__(self, model, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, cell_exit_extra_step=False,
64                  skip_target=0.4, temperature=None, branch_bias=0.25, entropy_reduction="sum"):
65         super().__init__(model)
66         self.lstm_size = lstm_size
67         self.lstm_num_layers = lstm_num_layers
68         self.tanh_constant = tanh_constant
69         self.temperature = temperature
70         self.cell_exit_extra_step = cell_exit_extra_step
71         self.skip_target = skip_target
72         self.branch_bias = branch_bias
73 
74         self.lstm = StackedLSTMCell(self.lstm_num_layers, self.lstm_size, False)
75         self.attn_anchor = nn.Linear(self.lstm_size, self.lstm_size, bias=False)
76         self.attn_query = nn.Linear(self.lstm_size, self.lstm_size, bias=False)
77         self.v_attn = nn.Linear(self.lstm_size, 1, bias=False)
78         self.g_emb = nn.Parameter(torch.randn(1, self.lstm_size) * 0.1)
79         self.skip_targets = nn.Parameter(torch.tensor([1.0 - self.skip_target, self.skip_target]), requires_grad=False)  # pylint: disable=not-callable
80         assert entropy_reduction in ["sum", "mean"], "Entropy reduction must be one of sum and mean."
81         self.entropy_reduction = torch.sum if entropy_reduction == "sum" else torch.mean
82         self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="none")
83         self.bias_dict = nn.ParameterDict()
84 
85         self.max_layer_choice = 0
86         for mutable in self.mutables:
87             if isinstance(mutable, LayerChoice):
88                 if self.max_layer_choice == 0:
89                     self.max_layer_choice = len(mutable)
90                 assert self.max_layer_choice == len(mutable), \
91                     "ENAS mutator requires all layer choice have the same number of candidates."
92                 # We are judging by keys and module types to add biases to layer choices. Needs refactor.
93                 if "reduce" in mutable.key:
94                     def is_conv(choice):
95                         return "conv" in str(type(choice)).lower()
96                     bias = torch.tensor([self.branch_bias if is_conv(choice) else -self.branch_bias  # pylint: disable=not-callable
97                                          for choice in mutable])
98                     self.bias_dict[mutable.key] = nn.Parameter(bias, requires_grad=False)
99 
100         self.embedding = nn.Embedding(self.max_layer_choice + 1, self.lstm_size)
101         self.soft = nn.Linear(self.lstm_size, self.max_layer_choice, bias=False)
102 
103     def sample_search(self):
104         self._initialize()
105         self._sample(self.mutables)
106         return self._choices
107 
108     def sample_final(self):
109         return self.sample_search()
110 
111     def _sample(self, tree):
112         mutable = tree.mutable
113         if isinstance(mutable, LayerChoice) and mutable.key not in self._choices:
114             self._choices[mutable.key] = self._sample_layer_choice(mutable)
115         elif isinstance(mutable, InputChoice) and mutable.key not in self._choices:
116             self._choices[mutable.key] = self._sample_input_choice(mutable)
117         for child in tree.children:
118             self._sample(child)
119         if isinstance(mutable, MutableScope) and mutable.key not in self._anchors_hid:
120             if self.cell_exit_extra_step:
121                 self._lstm_next_step()
122             self._mark_anchor(mutable.key)
123 
124     def _initialize(self):
125         self._choices = dict()
126         self._anchors_hid = dict()
127         self._inputs = self.g_emb.data
128         self._c = [torch.zeros((1, self.lstm_size),
129                                dtype=self._inputs.dtype,
130                                device=self._inputs.device) for _ in range(self.lstm_num_layers)]
131         self._h = [torch.zeros((1, self.lstm_size),
132                                dtype=self._inputs.dtype,
133                                device=self._inputs.device) for _ in range(self.lstm_num_layers)]
134         self.sample_log_prob = 0
135         self.sample_entropy = 0
136         self.sample_skip_penalty = 0
137 
138     def _lstm_next_step(self):
139         self._c, self._h = self.lstm(self._inputs, (self._c, self._h))
140 
141     def _mark_anchor(self, key):
142         self._anchors_hid[key] = self._h[-1]
143 
144     def _sample_layer_choice(self, mutable):
145         self._lstm_next_step()
146         logit = self.soft(self._h[-1])
147         if self.temperature is not None:
148             logit /= self.temperature
149         if self.tanh_constant is not None:
150             logit = self.tanh_constant * torch.tanh(logit)
151         if mutable.key in self.bias_dict:
152             logit += self.bias_dict[mutable.key]
153         branch_id = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)
154         log_prob = self.cross_entropy_loss(logit, branch_id)
155         self.sample_log_prob += self.entropy_reduction(log_prob)
156         entropy = (log_prob * torch.exp(-log_prob)).detach()  # pylint: disable=invalid-unary-operand-type
157         self.sample_entropy += self.entropy_reduction(entropy)
158         self._inputs = self.embedding(branch_id)
159         return F.one_hot(branch_id, num_classes=self.max_layer_choice).bool().view(-1)
160 
161     def _sample_input_choice(self, mutable):
162         query, anchors = [], []
163         for label in mutable.choose_from:
164             if label not in self._anchors_hid:
165                 self._lstm_next_step()
166                 self._mark_anchor(label)  # empty loop, fill not found
167             query.append(self.attn_anchor(self._anchors_hid[label]))
168             anchors.append(self._anchors_hid[label])
169         query = torch.cat(query, 0)
170         query = torch.tanh(query + self.attn_query(self._h[-1]))
171         query = self.v_attn(query)
172         if self.temperature is not None:
173             query /= self.temperature
174         if self.tanh_constant is not None:
175             query = self.tanh_constant * torch.tanh(query)
176 
177         if mutable.n_chosen is None:
178             logit = torch.cat([-query, query], 1)  # pylint: disable=invalid-unary-operand-type
179 
180             skip = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)
181             skip_prob = torch.sigmoid(logit)
182             kl = torch.sum(skip_prob * torch.log(skip_prob / self.skip_targets))
183             self.sample_skip_penalty += kl
184             log_prob = self.cross_entropy_loss(logit, skip)
185             self._inputs = (torch.matmul(skip.float(), torch.cat(anchors, 0)) / (1. + torch.sum(skip))).unsqueeze(0)
186         else:
187             assert mutable.n_chosen == 1, "Input choice must select exactly one or any in ENAS."
188             logit = query.view(1, -1)
189             index = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)
190             skip = F.one_hot(index, num_classes=mutable.n_candidates).view(-1)
191             log_prob = self.cross_entropy_loss(logit, index)
192             self._inputs = anchors[index.item()]
193 
194         self.sample_log_prob += self.entropy_reduction(log_prob)
195         entropy = (log_prob * torch.exp(-log_prob)).detach()  # pylint: disable=invalid-unary-operand-type
196         self.sample_entropy += self.entropy_reduction(entropy)
197         return skip.bool()
198
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py b/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py
--- a/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py
+++ b/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py
@@ -17,16 +17,16 @@
                                            for _ in range(self.lstm_num_layers)])
 
     def forward(self, inputs, hidden):
-        prev_c, prev_h = hidden
-        next_c, next_h = [], []
+        prev_h, prev_c = hidden
+        next_h, next_c = [], []
         for i, m in enumerate(self.lstm_modules):
-            curr_c, curr_h = m(inputs, (prev_c[i], prev_h[i]))
+            curr_h, curr_c = m(inputs, (prev_h[i], prev_c[i]))
             next_c.append(curr_c)
             next_h.append(curr_h)
             # current implementation only supports batch size equals 1,
             # but the algorithm does not necessarily have this limitation
             inputs = curr_h[-1].view(1, -1)
-        return next_c, next_h
+        return next_h, next_c
 
 
 class EnasMutator(Mutator):
@@ -136,7 +136,7 @@
         self.sample_skip_penalty = 0
 
     def _lstm_next_step(self):
-        self._c, self._h = self.lstm(self._inputs, (self._c, self._h))
+        self._h, self._c = self.lstm(self._inputs, (self._h, self._c))
 
     def _mark_anchor(self, key):
         self._anchors_hid[key] = self._h[-1]

{"golden_diff": "diff --git a/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py b/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py\n--- a/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py\n+++ b/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py\n@@ -17,16 +17,16 @@\n                                            for _ in range(self.lstm_num_layers)])\n \n     def forward(self, inputs, hidden):\n-        prev_c, prev_h = hidden\n-        next_c, next_h = [], []\n+        prev_h, prev_c = hidden\n+        next_h, next_c = [], []\n         for i, m in enumerate(self.lstm_modules):\n-            curr_c, curr_h = m(inputs, (prev_c[i], prev_h[i]))\n+            curr_h, curr_c = m(inputs, (prev_h[i], prev_c[i]))\n             next_c.append(curr_c)\n             next_h.append(curr_h)\n             # current implementation only supports batch size equals 1,\n             # but the algorithm does not necessarily have this limitation\n             inputs = curr_h[-1].view(1, -1)\n-        return next_c, next_h\n+        return next_h, next_c\n \n \n class EnasMutator(Mutator):\n@@ -136,7 +136,7 @@\n         self.sample_skip_penalty = 0\n \n     def _lstm_next_step(self):\n-        self._c, self._h = self.lstm(self._inputs, (self._c, self._h))\n+        self._h, self._c = self.lstm(self._inputs, (self._h, self._c))\n \n     def _mark_anchor(self, key):\n         self._anchors_hid[key] = self._h[-1]\n", "issue": "Is StackedLSTMCell correct implementation?\nFor enas, StackedLSTMCell is implemented as below\r\n https://github.com/microsoft/nni/blob/bf8be1e705cdc0c153d6f38d1a83bd3fe55d7a1e/src/sdk/pynni/nni/nas/pytorch/enas/mutator.py#L23-L28\r\n\r\nnn.LSTMCell outputs (hidden state, cell state) but the order of output is changed in the implementation.\r\nAs a result, the cell state is used for output of LSTM.\r\n\r\nIs this an intention or wrong implementation?\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# Licensed under the MIT license.\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom nni.nas.pytorch.mutator import Mutator\nfrom nni.nas.pytorch.mutables import LayerChoice, InputChoice, MutableScope\n\n\nclass StackedLSTMCell(nn.Module):\n    def __init__(self, layers, size, bias):\n        super().__init__()\n        self.lstm_num_layers = layers\n        self.lstm_modules = nn.ModuleList([nn.LSTMCell(size, size, bias=bias)\n                                           for _ in range(self.lstm_num_layers)])\n\n    def forward(self, inputs, hidden):\n        prev_c, prev_h = hidden\n        next_c, next_h = [], []\n        for i, m in enumerate(self.lstm_modules):\n            curr_c, curr_h = m(inputs, (prev_c[i], prev_h[i]))\n            next_c.append(curr_c)\n            next_h.append(curr_h)\n            # current implementation only supports batch size equals 1,\n            # but the algorithm does not necessarily have this limitation\n            inputs = curr_h[-1].view(1, -1)\n        return next_c, next_h\n\n\nclass EnasMutator(Mutator):\n    \"\"\"\n    A mutator that mutates the graph with RL.\n\n    Parameters\n    ----------\n    model : nn.Module\n        PyTorch model.\n    lstm_size : int\n        Controller LSTM hidden units.\n    lstm_num_layers : int\n        Number of layers for stacked LSTM.\n    tanh_constant : float\n        Logits will be equal to ``tanh_constant * tanh(logits)``. Don't use ``tanh`` if this value is ``None``.\n    cell_exit_extra_step : bool\n        If true, RL controller will perform an extra step at the exit of each MutableScope, dump the hidden state\n        and mark it as the hidden state of this MutableScope. This is to align with the original implementation of paper.\n    skip_target : float\n        Target probability that skipconnect will appear.\n    temperature : float\n        Temperature constant that divides the logits.\n    branch_bias : float\n        Manual bias applied to make some operations more likely to be chosen.\n        Currently this is implemented with a hardcoded match rule that aligns with original repo.\n        If a mutable has a ``reduce`` in its key, all its op choices\n        that contains `conv` in their typename will receive a bias of ``+self.branch_bias`` initially; while others\n        receive a bias of ``-self.branch_bias``.\n    entropy_reduction : str\n        Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced.\n    \"\"\"\n\n    def __init__(self, model, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, cell_exit_extra_step=False,\n                 skip_target=0.4, temperature=None, branch_bias=0.25, entropy_reduction=\"sum\"):\n        super().__init__(model)\n        self.lstm_size = lstm_size\n        self.lstm_num_layers = lstm_num_layers\n        self.tanh_constant = tanh_constant\n        self.temperature = temperature\n        self.cell_exit_extra_step = cell_exit_extra_step\n        self.skip_target = skip_target\n        self.branch_bias = branch_bias\n\n        self.lstm = StackedLSTMCell(self.lstm_num_layers, self.lstm_size, False)\n        self.attn_anchor = nn.Linear(self.lstm_size, self.lstm_size, bias=False)\n        self.attn_query = nn.Linear(self.lstm_size, self.lstm_size, bias=False)\n        self.v_attn = nn.Linear(self.lstm_size, 1, bias=False)\n        self.g_emb = nn.Parameter(torch.randn(1, self.lstm_size) * 0.1)\n        self.skip_targets = nn.Parameter(torch.tensor([1.0 - self.skip_target, self.skip_target]), requires_grad=False)  # pylint: disable=not-callable\n        assert entropy_reduction in [\"sum\", \"mean\"], \"Entropy reduction must be one of sum and mean.\"\n        self.entropy_reduction = torch.sum if entropy_reduction == \"sum\" else torch.mean\n        self.cross_entropy_loss = nn.CrossEntropyLoss(reduction=\"none\")\n        self.bias_dict = nn.ParameterDict()\n\n        self.max_layer_choice = 0\n        for mutable in self.mutables:\n            if isinstance(mutable, LayerChoice):\n                if self.max_layer_choice == 0:\n                    self.max_layer_choice = len(mutable)\n                assert self.max_layer_choice == len(mutable), \\\n                    \"ENAS mutator requires all layer choice have the same number of candidates.\"\n                # We are judging by keys and module types to add biases to layer choices. Needs refactor.\n                if \"reduce\" in mutable.key:\n                    def is_conv(choice):\n                        return \"conv\" in str(type(choice)).lower()\n                    bias = torch.tensor([self.branch_bias if is_conv(choice) else -self.branch_bias  # pylint: disable=not-callable\n                                         for choice in mutable])\n                    self.bias_dict[mutable.key] = nn.Parameter(bias, requires_grad=False)\n\n        self.embedding = nn.Embedding(self.max_layer_choice + 1, self.lstm_size)\n        self.soft = nn.Linear(self.lstm_size, self.max_layer_choice, bias=False)\n\n    def sample_search(self):\n        self._initialize()\n        self._sample(self.mutables)\n        return self._choices\n\n    def sample_final(self):\n        return self.sample_search()\n\n    def _sample(self, tree):\n        mutable = tree.mutable\n        if isinstance(mutable, LayerChoice) and mutable.key not in self._choices:\n            self._choices[mutable.key] = self._sample_layer_choice(mutable)\n        elif isinstance(mutable, InputChoice) and mutable.key not in self._choices:\n            self._choices[mutable.key] = self._sample_input_choice(mutable)\n        for child in tree.children:\n            self._sample(child)\n        if isinstance(mutable, MutableScope) and mutable.key not in self._anchors_hid:\n            if self.cell_exit_extra_step:\n                self._lstm_next_step()\n            self._mark_anchor(mutable.key)\n\n    def _initialize(self):\n        self._choices = dict()\n        self._anchors_hid = dict()\n        self._inputs = self.g_emb.data\n        self._c = [torch.zeros((1, self.lstm_size),\n                               dtype=self._inputs.dtype,\n                               device=self._inputs.device) for _ in range(self.lstm_num_layers)]\n        self._h = [torch.zeros((1, self.lstm_size),\n                               dtype=self._inputs.dtype,\n                               device=self._inputs.device) for _ in range(self.lstm_num_layers)]\n        self.sample_log_prob = 0\n        self.sample_entropy = 0\n        self.sample_skip_penalty = 0\n\n    def _lstm_next_step(self):\n        self._c, self._h = self.lstm(self._inputs, (self._c, self._h))\n\n    def _mark_anchor(self, key):\n        self._anchors_hid[key] = self._h[-1]\n\n    def _sample_layer_choice(self, mutable):\n        self._lstm_next_step()\n        logit = self.soft(self._h[-1])\n        if self.temperature is not None:\n            logit /= self.temperature\n        if self.tanh_constant is not None:\n            logit = self.tanh_constant * torch.tanh(logit)\n        if mutable.key in self.bias_dict:\n            logit += self.bias_dict[mutable.key]\n        branch_id = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)\n        log_prob = self.cross_entropy_loss(logit, branch_id)\n        self.sample_log_prob += self.entropy_reduction(log_prob)\n        entropy = (log_prob * torch.exp(-log_prob)).detach()  # pylint: disable=invalid-unary-operand-type\n        self.sample_entropy += self.entropy_reduction(entropy)\n        self._inputs = self.embedding(branch_id)\n        return F.one_hot(branch_id, num_classes=self.max_layer_choice).bool().view(-1)\n\n    def _sample_input_choice(self, mutable):\n        query, anchors = [], []\n        for label in mutable.choose_from:\n            if label not in self._anchors_hid:\n                self._lstm_next_step()\n                self._mark_anchor(label)  # empty loop, fill not found\n            query.append(self.attn_anchor(self._anchors_hid[label]))\n            anchors.append(self._anchors_hid[label])\n        query = torch.cat(query, 0)\n        query = torch.tanh(query + self.attn_query(self._h[-1]))\n        query = self.v_attn(query)\n        if self.temperature is not None:\n            query /= self.temperature\n        if self.tanh_constant is not None:\n            query = self.tanh_constant * torch.tanh(query)\n\n        if mutable.n_chosen is None:\n            logit = torch.cat([-query, query], 1)  # pylint: disable=invalid-unary-operand-type\n\n            skip = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)\n            skip_prob = torch.sigmoid(logit)\n            kl = torch.sum(skip_prob * torch.log(skip_prob / self.skip_targets))\n            self.sample_skip_penalty += kl\n            log_prob = self.cross_entropy_loss(logit, skip)\n            self._inputs = (torch.matmul(skip.float(), torch.cat(anchors, 0)) / (1. + torch.sum(skip))).unsqueeze(0)\n        else:\n            assert mutable.n_chosen == 1, \"Input choice must select exactly one or any in ENAS.\"\n            logit = query.view(1, -1)\n            index = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)\n            skip = F.one_hot(index, num_classes=mutable.n_candidates).view(-1)\n            log_prob = self.cross_entropy_loss(logit, index)\n            self._inputs = anchors[index.item()]\n\n        self.sample_log_prob += self.entropy_reduction(log_prob)\n        entropy = (log_prob * torch.exp(-log_prob)).detach()  # pylint: disable=invalid-unary-operand-type\n        self.sample_entropy += self.entropy_reduction(entropy)\n        return skip.bool()\n", "path": "src/sdk/pynni/nni/nas/pytorch/enas/mutator.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# Licensed under the MIT license.\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom nni.nas.pytorch.mutator import Mutator\nfrom nni.nas.pytorch.mutables import LayerChoice, InputChoice, MutableScope\n\n\nclass StackedLSTMCell(nn.Module):\n    def __init__(self, layers, size, bias):\n        super().__init__()\n        self.lstm_num_layers = layers\n        self.lstm_modules = nn.ModuleList([nn.LSTMCell(size, size, bias=bias)\n                                           for _ in range(self.lstm_num_layers)])\n\n    def forward(self, inputs, hidden):\n        prev_h, prev_c = hidden\n        next_h, next_c = [], []\n        for i, m in enumerate(self.lstm_modules):\n            curr_h, curr_c = m(inputs, (prev_h[i], prev_c[i]))\n            next_c.append(curr_c)\n            next_h.append(curr_h)\n            # current implementation only supports batch size equals 1,\n            # but the algorithm does not necessarily have this limitation\n            inputs = curr_h[-1].view(1, -1)\n        return next_h, next_c\n\n\nclass EnasMutator(Mutator):\n    \"\"\"\n    A mutator that mutates the graph with RL.\n\n    Parameters\n    ----------\n    model : nn.Module\n        PyTorch model.\n    lstm_size : int\n        Controller LSTM hidden units.\n    lstm_num_layers : int\n        Number of layers for stacked LSTM.\n    tanh_constant : float\n        Logits will be equal to ``tanh_constant * tanh(logits)``. Don't use ``tanh`` if this value is ``None``.\n    cell_exit_extra_step : bool\n        If true, RL controller will perform an extra step at the exit of each MutableScope, dump the hidden state\n        and mark it as the hidden state of this MutableScope. This is to align with the original implementation of paper.\n    skip_target : float\n        Target probability that skipconnect will appear.\n    temperature : float\n        Temperature constant that divides the logits.\n    branch_bias : float\n        Manual bias applied to make some operations more likely to be chosen.\n        Currently this is implemented with a hardcoded match rule that aligns with original repo.\n        If a mutable has a ``reduce`` in its key, all its op choices\n        that contains `conv` in their typename will receive a bias of ``+self.branch_bias`` initially; while others\n        receive a bias of ``-self.branch_bias``.\n    entropy_reduction : str\n        Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced.\n    \"\"\"\n\n    def __init__(self, model, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, cell_exit_extra_step=False,\n                 skip_target=0.4, temperature=None, branch_bias=0.25, entropy_reduction=\"sum\"):\n        super().__init__(model)\n        self.lstm_size = lstm_size\n        self.lstm_num_layers = lstm_num_layers\n        self.tanh_constant = tanh_constant\n        self.temperature = temperature\n        self.cell_exit_extra_step = cell_exit_extra_step\n        self.skip_target = skip_target\n        self.branch_bias = branch_bias\n\n        self.lstm = StackedLSTMCell(self.lstm_num_layers, self.lstm_size, False)\n        self.attn_anchor = nn.Linear(self.lstm_size, self.lstm_size, bias=False)\n        self.attn_query = nn.Linear(self.lstm_size, self.lstm_size, bias=False)\n        self.v_attn = nn.Linear(self.lstm_size, 1, bias=False)\n        self.g_emb = nn.Parameter(torch.randn(1, self.lstm_size) * 0.1)\n        self.skip_targets = nn.Parameter(torch.tensor([1.0 - self.skip_target, self.skip_target]), requires_grad=False)  # pylint: disable=not-callable\n        assert entropy_reduction in [\"sum\", \"mean\"], \"Entropy reduction must be one of sum and mean.\"\n        self.entropy_reduction = torch.sum if entropy_reduction == \"sum\" else torch.mean\n        self.cross_entropy_loss = nn.CrossEntropyLoss(reduction=\"none\")\n        self.bias_dict = nn.ParameterDict()\n\n        self.max_layer_choice = 0\n        for mutable in self.mutables:\n            if isinstance(mutable, LayerChoice):\n                if self.max_layer_choice == 0:\n                    self.max_layer_choice = len(mutable)\n                assert self.max_layer_choice == len(mutable), \\\n                    \"ENAS mutator requires all layer choice have the same number of candidates.\"\n                # We are judging by keys and module types to add biases to layer choices. Needs refactor.\n                if \"reduce\" in mutable.key:\n                    def is_conv(choice):\n                        return \"conv\" in str(type(choice)).lower()\n                    bias = torch.tensor([self.branch_bias if is_conv(choice) else -self.branch_bias  # pylint: disable=not-callable\n                                         for choice in mutable])\n                    self.bias_dict[mutable.key] = nn.Parameter(bias, requires_grad=False)\n\n        self.embedding = nn.Embedding(self.max_layer_choice + 1, self.lstm_size)\n        self.soft = nn.Linear(self.lstm_size, self.max_layer_choice, bias=False)\n\n    def sample_search(self):\n        self._initialize()\n        self._sample(self.mutables)\n        return self._choices\n\n    def sample_final(self):\n        return self.sample_search()\n\n    def _sample(self, tree):\n        mutable = tree.mutable\n        if isinstance(mutable, LayerChoice) and mutable.key not in self._choices:\n            self._choices[mutable.key] = self._sample_layer_choice(mutable)\n        elif isinstance(mutable, InputChoice) and mutable.key not in self._choices:\n            self._choices[mutable.key] = self._sample_input_choice(mutable)\n        for child in tree.children:\n            self._sample(child)\n        if isinstance(mutable, MutableScope) and mutable.key not in self._anchors_hid:\n            if self.cell_exit_extra_step:\n                self._lstm_next_step()\n            self._mark_anchor(mutable.key)\n\n    def _initialize(self):\n        self._choices = dict()\n        self._anchors_hid = dict()\n        self._inputs = self.g_emb.data\n        self._c = [torch.zeros((1, self.lstm_size),\n                               dtype=self._inputs.dtype,\n                               device=self._inputs.device) for _ in range(self.lstm_num_layers)]\n        self._h = [torch.zeros((1, self.lstm_size),\n                               dtype=self._inputs.dtype,\n                               device=self._inputs.device) for _ in range(self.lstm_num_layers)]\n        self.sample_log_prob = 0\n        self.sample_entropy = 0\n        self.sample_skip_penalty = 0\n\n    def _lstm_next_step(self):\n        self._h, self._c = self.lstm(self._inputs, (self._h, self._c))\n\n    def _mark_anchor(self, key):\n        self._anchors_hid[key] = self._h[-1]\n\n    def _sample_layer_choice(self, mutable):\n        self._lstm_next_step()\n        logit = self.soft(self._h[-1])\n        if self.temperature is not None:\n            logit /= self.temperature\n        if self.tanh_constant is not None:\n            logit = self.tanh_constant * torch.tanh(logit)\n        if mutable.key in self.bias_dict:\n            logit += self.bias_dict[mutable.key]\n        branch_id = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)\n        log_prob = self.cross_entropy_loss(logit, branch_id)\n        self.sample_log_prob += self.entropy_reduction(log_prob)\n        entropy = (log_prob * torch.exp(-log_prob)).detach()  # pylint: disable=invalid-unary-operand-type\n        self.sample_entropy += self.entropy_reduction(entropy)\n        self._inputs = self.embedding(branch_id)\n        return F.one_hot(branch_id, num_classes=self.max_layer_choice).bool().view(-1)\n\n    def _sample_input_choice(self, mutable):\n        query, anchors = [], []\n        for label in mutable.choose_from:\n            if label not in self._anchors_hid:\n                self._lstm_next_step()\n                self._mark_anchor(label)  # empty loop, fill not found\n            query.append(self.attn_anchor(self._anchors_hid[label]))\n            anchors.append(self._anchors_hid[label])\n        query = torch.cat(query, 0)\n        query = torch.tanh(query + self.attn_query(self._h[-1]))\n        query = self.v_attn(query)\n        if self.temperature is not None:\n            query /= self.temperature\n        if self.tanh_constant is not None:\n            query = self.tanh_constant * torch.tanh(query)\n\n        if mutable.n_chosen is None:\n            logit = torch.cat([-query, query], 1)  # pylint: disable=invalid-unary-operand-type\n\n            skip = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)\n            skip_prob = torch.sigmoid(logit)\n            kl = torch.sum(skip_prob * torch.log(skip_prob / self.skip_targets))\n            self.sample_skip_penalty += kl\n            log_prob = self.cross_entropy_loss(logit, skip)\n            self._inputs = (torch.matmul(skip.float(), torch.cat(anchors, 0)) / (1. + torch.sum(skip))).unsqueeze(0)\n        else:\n            assert mutable.n_chosen == 1, \"Input choice must select exactly one or any in ENAS.\"\n            logit = query.view(1, -1)\n            index = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1)\n            skip = F.one_hot(index, num_classes=mutable.n_candidates).view(-1)\n            log_prob = self.cross_entropy_loss(logit, index)\n            self._inputs = anchors[index.item()]\n\n        self.sample_log_prob += self.entropy_reduction(log_prob)\n        entropy = (log_prob * torch.exp(-log_prob)).detach()  # pylint: disable=invalid-unary-operand-type\n        self.sample_entropy += self.entropy_reduction(entropy)\n        return skip.bool()\n", "path": "src/sdk/pynni/nni/nas/pytorch/enas/mutator.py"}]}