cleanup search beam logic

#18

Author: albertz

nn/search.py

@@ -11,14 +11,6 @@ class SearchFuncInterface:
   Interface for search.
   """
 
-  def get_beam(self) -> nn.SearchBeam:
-    """
-    Return a new beam instance. dependency and name is still unset and will be set outside.
-    This overwrites whatever is returned by :func:`choice`,
-    and copy_as_prev_frame() is used to set the initial (prev) state beam.
-    """
-    raise NotImplementedError
-
   def choice(self, *, probs: nn.Tensor, probs_type: str) -> nn.Tensor:
     """
     Given an output tensor (logits or log prop), returns a beam of chosen indices.
@@ -43,10 +35,6 @@ def __init__(self, beam_size: int, max_seq_len: nn.Tensor):
     self.beam_size = beam_size
     self.max_seq_len = max_seq_len
 
-  def get_beam(self):
-    """beam"""
-    return nn.SearchBeam(beam_size=self.beam_size)
-
   def choice(self, *, probs: nn.Tensor, probs_type: str) -> nn.Tensor:
     """nn.choice"""
     return nn.choice(

nn/transformer.py

@@ -380,13 +380,6 @@ def __call__(self, source: nn.Tensor, *,
       assert target_spatial_axis, f"{self}: Target spatial axis must be specified when target is given"
     loop = nn.Loop(axis=target_spatial_axis)
     loop.state = state if state else self.default_initial_state()
-    beam = None
-    if search:
-      beam = search.get_beam()
-      beam.name = f"{nn.NameCtx.current_ctx().get_abs_name()}/target"
-      beam.dependency = beam.copy_as_prev_frame()
-      for x in loop.state.deep_tensors():
-        x.data.beam = beam.dependency
     with loop:
       prev_target_embed = self.target_embedding(loop.state.target)
       output, loop.state.decoder = self.decoder(
@@ -397,7 +390,6 @@ def __call__(self, source: nn.Tensor, *,
       if search:
         search.apply_loop(loop)
         choice = search.choice(probs=logits, probs_type="logits")
-        choice.data.beam = beam
         loop.state.target = choice
         loop.end(loop.state.target == self.target_eos_symbol, include_eos=False)
       else: