Do not sort BIAS_STATE in beam search (#971)

CHANGELOG.md

@@ -11,6 +11,11 @@ Note that Sockeye has checks in place to not translate with an old model that wa
 
 Each version section may have have subsections for: _Added_, _Changed_, _Removed_, _Deprecated_, and _Fixed_.
 
+## [2.3.23]
+### Changed
+
+- Do not sort BIAS_STATE in beam search. It is constant across decoder steps.
+
 ## [2.3.22]
 ### Fixed
 

sockeye/__init__.py

@@ -11,4 +11,4 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-__version__ = '2.3.22'
+__version__ = '2.3.23'

sockeye/beam_search.py

@@ -471,13 +471,13 @@ def hybrid_forward(self, F, best_hyp_indices, *states):
         sorted_states = []
         assert len(states) == len(self.flat_structure), "Number of states do not match the defined state structure"
         for state, state_format in zip(states, self.flat_structure):
-            if state_format == C.STEP_STATE or state_format == C.BIAS_STATE:
-                # Steps and source_bias have batch dimension on axis 0
+            if state_format == C.STEP_STATE:
+                # Step has batch dimension on axis 0
                 sorted_state = F.take(state, best_hyp_indices)
             elif state_format == C.DECODER_STATE:
                 # Decoder and encoder layer states have batch dimension on axis 1
                 sorted_state = F.take(state, best_hyp_indices, axis=1)
-            elif state_format == C.ENCODER_STATE:
+            elif state_format == C.ENCODER_STATE or state_format == C.BIAS_STATE:
                 # No need for takes on encoder layer states
                 sorted_state = state
             else: