class_docstr: Beam Search implementation ported from ESPNet implementation -
  https://github.com/espnet/espnet/blob/master/espnet/nets/beam_search_transducer.py
  
  Sequence level beam decoding or batched-beam decoding, performed auto-repressively
  depending on the search type chosen.
  
  Args:
      decoder_model: rnnt_utils.AbstractRNNTDecoder implementation.
      joint_model: rnnt_utils.AbstractRNNTJoint implementation.
  
      beam_size: number of beams for beam search. Must be a positive integer >= 1.
          If beam size is 1, defaults to stateful greedy search.
          This greedy search might result in slightly different results than
          the greedy results obtained by GreedyRNNTInfer due to implementation differences.
  
          For accurate greedy results, please use GreedyRNNTInfer or GreedyBatchedRNNTInfer.
  
      search_type: str representing the type of beam search to perform.
          Must be one of ['beam', 'tsd', 'alsd']. 'nsc' is currently not supported.
  
          Algoritm used:
          `beam` - basic beam search strategy. Larger beams generally result in better decoding,
              however the time required for the search also grows steadily.
  
          `tsd` - time synchronous decoding. Please refer to the paper:
              [Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040)
              for details on the algorithm implemented.
  
              Time synchronous decoding (TSD) execution time grows by the factor T * max_symmetric_expansions.
              For longer sequences, T is greater, and can therefore take a long time for beams to obtain
              good results. This also requires greater memory to execute.
  
          `alsd` - alignment-length synchronous decoding. Please refer to the paper:
              [Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040)
              for details on the algorithm implemented.
  
              Alignment-length synchronous decoding (ALSD) execution time is faster than TSD, with growth
              factor of T + U_max, where U_max is the maximum target length expected during execution.
  
              Generally, T + U_max < T * max_symmetric_expansions. However, ALSD beams are non-unique,
              therefore it is required to use larger beam sizes to achieve the same (or close to the same)
              decoding accuracy as TSD.
  
              For a given decoding accuracy, it is possible to attain faster decoding via ALSD than TSD.
  
          `maes` = modified adaptive expansion searcn. Please refer to the paper:
              [Accelerating RNN Transducer Inference via Adaptive Expansion Search](https://ieeexplore.ieee.org/document/9250505)
  
              Modified Adaptive Synchronous Decoding (mAES) execution time is adaptive w.r.t the
              number of expansions (for tokens) required per timestep. The number of expansions can usually
              be constrained to 1 or 2, and in most cases 2 is sufficient.
  
              This beam search technique can possibly obtain superior WER while sacrificing some evaluation time.
  
      score_norm: bool, whether to normalize the scores of the log probabilities.
  
      return_best_hypothesis: bool, decides whether to return a single hypothesis (the best out of N),
          or return all N hypothesis (sorted with best score first). The container class changes based
          this flag -
          When set to True (default), returns a single Hypothesis.
          When set to False, returns a NBestHypotheses container, which contains a list of Hypothesis.
  
      # The following arguments are specific to the chosen `search_type`
  
      tsd_max_sym_exp_per_step: Used for `search_type=tsd`. The maximum symmetric expansions allowed
          per timestep during beam search. Larger values should be used to attempt decoding of longer
          sequences, but this in turn increases execution time and memory usage.
  
      alsd_max_target_len: Used for `search_type=alsd`. The maximum expected target sequence length
          during beam search. Larger values allow decoding of longer sequences at the expense of
          execution time and memory.
  
      # The following two flags are placeholders and unused until `nsc` implementation is stabilized.
      nsc_max_timesteps_expansion: Unused int.
  
      nsc_prefix_alpha: Unused int.
  
      # mAES flags
      maes_num_steps: Number of adaptive steps to take. From the paper, 2 steps is generally sufficient. int > 1.
  
      maes_prefix_alpha: Maximum prefix length in prefix search. Must be an integer, and is advised to keep this as 1
          in order to reduce expensive beam search cost later. int >= 0.
  
      maes_expansion_beta: Maximum number of prefix expansions allowed, in addition to the beam size.
          Effectively, the number of hypothesis = beam_size + maes_expansion_beta. Must be an int >= 0,
          and affects the speed of inference since large values will perform large beam search in the next step.
  
      maes_expansion_gamma: Float pruning threshold used in the prune-by-value step when computing the expansions.
          The default (2.3) is selected from the paper. It performs a comparison (max_log_prob - gamma <= log_prob[v])
          where v is all vocabulary indices in the Vocab set and max_log_prob is the "most" likely token to be
          predicted. Gamma therefore provides a margin of additional tokens which can be potential candidates for
          expansion apart from the "most likely" candidate.
          Lower values will reduce the number of expansions (by increasing pruning-by-value, thereby improving speed
          but hurting accuracy). Higher values will increase the number of expansions (by reducing pruning-by-value,
          thereby reducing speed but potentially improving accuracy). This is a hyper parameter to be experimentally
          tuned on a validation set.
  
      softmax_temperature: Scales the logits of the joint prior to computing log_softmax.
  
      preserve_alignments: Bool flag which preserves the history of alignments generated during
          beam decoding (sample). When set to true, the Hypothesis will contain
          the non-null value for `alignments` in it. Here, `alignments` is a List of List of Tensor (of length V + 1).
  
          The length of the list corresponds to the Acoustic Length (T).
          Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
          U is the number of target tokens for the current timestep Ti.
  
          NOTE: `preserve_alignments` is an invalid argument for any `search_type`
          other than basic beam search.
  
      ngram_lm_model: str
          The path to the N-gram LM
      ngram_lm_alpha: float
          Alpha weight of N-gram LM
      tokens_type: str
          Tokenization type ['subword', 'char']
docstr: Alignment-length synchronous beam search implementation.
  Based on https://ieeexplore.ieee.org/document/9053040
  
  Args:
      h: Encoded speech features (1, T_max, D_enc)
  
  Returns:
      nbest_hyps: N-best decoding results
docstr_tok: ["Alignment-length","synchronous","beam","search","implementation",".","Based","on","https",":","\/\/ieeexplore.ieee.org\/document\/9053040","Args",":","h",":","Encoded","speech","features","(","1",",","T_max",",","D_enc",")","Returns",":","nbest_hyps",":","N-best","decoding","results"]
code: def align_length_sync_decoding(
      self,
      h: torch.Tensor,
      encoded_lengths: torch.Tensor,
      partial_hypotheses: Optional[Hypothesis] = None,
  ) -> List[Hypothesis]:
      """Alignment-length synchronous beam search implementation.
      Based on https://ieeexplore.ieee.org/document/9053040
  
      Args:
          h: Encoded speech features (1, T_max, D_enc)
  
      Returns:
          nbest_hyps: N-best decoding results
      """
      # delay this import here instead of at the beginning to avoid circular imports.
      from nemo.collections.asr.modules.rnnt import (
          RNNTDecoder,
          StatelessTransducerDecoder,
      )
  
      if partial_hypotheses is not None:
          raise NotImplementedError(
              "`partial_hypotheses` support is not supported"
          )
  
      # Precompute some constants for blank position
      ids = list(range(self.vocab_size + 1))
      ids.remove(self.blank)
  
      # Used when blank token is first vs last token
      if self.blank == 0:
          index_incr = 1
      else:
          index_incr = 0
  
      # prepare the batched beam states
      beam = min(self.beam_size, self.vocab_size)
  
      h = h[0]  # [T, D]
      h_length = int(encoded_lengths)
      beam_state = self.decoder.initialize_state(
          torch.zeros(beam, device=h.device, dtype=h.dtype)
      )  # [L, B, H], [L, B, H] for LSTMS
  
      # compute u_max as either a specific static limit,
      # or a multiple of current `h_length` dynamically.
      if type(self.alsd_max_target_length) == float:
          u_max = int(self.alsd_max_target_length * h_length)
      else:
          u_max = int(self.alsd_max_target_length)
  
      # Initialize first hypothesis for the beam (blank)
      B = [
          Hypothesis(
              y_sequence=[self.blank],
              score=0.0,
              dec_state=self.decoder.batch_select_state(beam_state, 0),
              timestep=[-1],
              length=0,
          )
      ]
  
      # Initialize alignments
      if self.preserve_alignments:
          B[0].alignments = [[]]
  
      final = []
      cache = {}
  
      # ALSD runs for T + U_max steps
      for i in range(h_length + u_max):
          # Update caches
          A = []
          B_ = []
          h_states = []
  
          # preserve the list of batch indices which are added into the list
          # and those which are removed from the list
          # This is necessary to perform state updates in the correct batch indices later
          batch_ids = list(
              range(len(B))
          )  # initialize as a list of all batch ids
          batch_removal_ids = (
              []
          )  # update with sample ids which are removed
  
          for bid, hyp in enumerate(B):
              u = len(hyp.y_sequence) - 1
              t = i - u
  
              if t > (h_length - 1):
                  batch_removal_ids.append(bid)
                  continue
  
              B_.append(hyp)
              h_states.append((t, h[t]))
  
          if B_:
              # Compute the subset of batch ids which were *not* removed from the list above
              sub_batch_ids = None
              if len(B_) != beam:
                  sub_batch_ids = batch_ids
                  for id in batch_removal_ids:
                      # sub_batch_ids contains list of ids *that were not removed*
                      sub_batch_ids.remove(id)
  
                  # extract the states of the sub batch only.
                  if isinstance(self.decoder, RNNTDecoder):
                      # LSTM decoder, state is [layer x batch x hidden]
                      beam_state_ = [
                          beam_state[state_id][:, sub_batch_ids, :]
                          for state_id in range(len(beam_state))
                      ]
                  elif isinstance(
                      self.decoder, StatelessTransducerDecoder
                  ):
                      # stateless decoder, state is [batch x hidden]
                      beam_state_ = [
                          beam_state[state_id][sub_batch_ids, :]
                          for state_id in range(len(beam_state))
                      ]
                  else:
                      raise NotImplementedError("Unknown decoder type.")
  
              else:
                  # If entire batch was used (none were removed), simply take all the states
                  beam_state_ = beam_state
  
              # Decode a batch/sub-batch of beam states and scores
              (
                  beam_y,
                  beam_state_,
                  beam_lm_tokens,
              ) = self.decoder.batch_score_hypothesis(
                  B_, cache, beam_state_
              )
  
              # If only a subset of batch ids were updated (some were removed)
              if sub_batch_ids is not None:
                  # For each state in the RNN (2 for LSTM)
                  for state_id in range(len(beam_state)):
                      # Update the current batch states with the sub-batch states (in the correct indices)
                      # These indices are specified by sub_batch_ids, the ids of samples which were updated.
                      if isinstance(self.decoder, RNNTDecoder):
                          # LSTM decoder, state is [layer x batch x hidden]
                          beam_state[state_id][
                              :, sub_batch_ids, :
                          ] = beam_state_[state_id][...]
                      elif isinstance(
                          self.decoder, StatelessTransducerDecoder
                      ):
                          # stateless decoder, state is [batch x hidden]
                          beam_state[state_id][
                              sub_batch_ids, :
                          ] = beam_state_[state_id][...]
                      else:
                          raise NotImplementedError(
                              "Unknown decoder type."
                          )
              else:
                  # If entire batch was updated, simply update all the states
                  beam_state = beam_state_
  
              # h_states = list of [t, h[t]]
              # so h[1] here is a h[t] of shape [D]
              # Simply stack all of the h[t] within the sub_batch/batch (T <= beam)
              h_enc = torch.stack(
                  [h[1] for h in h_states]
              )  # [T=beam, D]
              h_enc = h_enc.unsqueeze(
                  1
              )  # [B=beam, T=1, D]; batch over the beams
  
              # Extract the log probabilities and the predicted tokens
              beam_logp = torch.log_softmax(
                  self.joint.joint(h_enc, beam_y)
                  / self.softmax_temperature,
                  dim=-1,
              )  # [B=beam, 1, 1, V + 1]
              beam_logp = beam_logp[:, 0, 0, :]  # [B=beam, V + 1]
              beam_topk = beam_logp[:, ids].topk(beam, dim=-1)
  
              for j, hyp in enumerate(B_):
                  # For all updated samples in the batch, add it as the blank token
                  # In this step, we dont add a token but simply update score
                  new_hyp = Hypothesis(
                      score=(
                          hyp.score + float(beam_logp[j, self.blank])
                      ),
                      y_sequence=hyp.y_sequence[:],
                      dec_state=hyp.dec_state,
                      lm_state=hyp.lm_state,
                      timestep=hyp.timestep[:],
                      length=i,
                  )
  
                  if self.preserve_alignments:
                      new_hyp.alignments = copy.deepcopy(hyp.alignments)
  
                      # Add the alignment of blank at this step
                      new_hyp.alignments[-1].append(
                          (
                              beam_logp[j].clone().cpu(),
                              torch.tensor(
                                  self.blank, dtype=torch.int32
                              ),
                          )
                      )
  
                  # Add blank prediction to A
                  A.append(new_hyp)
  
                  # If the prediction "timestep" t has reached the length of the input sequence
                  # we can add it to the "finished" hypothesis list.
                  if h_states[j][0] == (h_length - 1):
                      final.append(new_hyp)
  
                  # Here, we carefully select the indices of the states that we want to preserve
                  # for the next token (non-blank) update.
                  if sub_batch_ids is not None:
                      h_states_idx = sub_batch_ids[j]
                  else:
                      h_states_idx = j
  
                  # for each current hypothesis j
                  # extract the top token score and top token id for the jth hypothesis
                  for logp, k in zip(
                      beam_topk[0][j], beam_topk[1][j] + index_incr
                  ):
                      # create new hypothesis and store in A
                      # Note: This loop does *not* include the blank token!
                      new_hyp = Hypothesis(
                          score=(hyp.score + float(logp)),
                          y_sequence=(hyp.y_sequence[:] + [int(k)]),
                          dec_state=self.decoder.batch_select_state(
                              beam_state, h_states_idx
                          ),
                          lm_state=hyp.lm_state,
                          timestep=hyp.timestep[:] + [i],
                          length=i,
                      )
  
                      if self.preserve_alignments:
                          new_hyp.alignments = copy.deepcopy(
                              hyp.alignments
                          )
  
                          # Add the alignment of Uj for this beam candidate at this step
                          new_hyp.alignments[-1].append(
                              (
                                  beam_logp[j].clone().cpu(),
                                  torch.tensor(
                                      new_hyp.y_sequence[-1],
                                      dtype=torch.int32,
                                  ),
                              )
                          )
  
                      A.append(new_hyp)
  
              # Prune and recombine same hypothesis
              # This may cause next beam to be smaller than max beam size
              # Therefore larger beam sizes may be required for better decoding.
              B = sorted(A, key=lambda x: x.score, reverse=True)[:beam]
              B = self.recombine_hypotheses(B)
  
              if self.preserve_alignments:
                  # convert Ti-th logits into a torch array
                  for B_i in B:
                      # Check if the last token emitted at last timestep was a blank
                      # If so, move to next timestep
                      logp, label = B_i.alignments[-1][
                          -1
                      ]  # The last alignment of this step
                      if int(label) == self.blank:
                          B_i.alignments.append(
                              []
                          )  # blank buffer for next timestep
  
          # If B_ is empty list, then we may be able to early exit
          elif len(batch_ids) == len(batch_removal_ids):
              # break early
              break
  
      if final:
          # Remove trailing empty list of alignments
          if self.preserve_alignments:
              for h in final:
                  if len(h.alignments[-1]) == 0:
                      del h.alignments[-1]
  
          return self.sort_nbest(final)
      else:
          # Remove trailing empty list of alignments
          if self.preserve_alignments:
              for h in B:
                  if len(h.alignments[-1]) == 0:
                      del h.alignments[-1]
  
          return B
  
imports: import copy
  from dataclasses import dataclass
  from typing import Any, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  from tqdm import tqdm
  from nemo.collections.asr.modules import rnnt_abstract
  from nemo.collections.asr.parts.utils.rnnt_utils import HATJointOutput, Hypothesis, NBestHypotheses, is_prefix, select_k_expansions
  from nemo.core.classes import Typing, typecheck
  from nemo.core.neural_types import AcousticEncodedRepresentation, HypothesisType, LengthsType, NeuralType
  from nemo.utils import logging
label_desc: Use image node_id 9 for calling the BeamRNNTInfer obj's underlying member method code with example usage: obj.align_length_sync_decoding(h, encoded_lengths, partial_hypotheses) and returns: self, B