text: # Copyright      2020  Mobvoi Inc.        (authors: Fangjun Kuang)
  #                      Xiaomi Corp.       (authors: Daniel Povey
  #                                                   Haowen Qiu)
  # See ../../../LICENSE for clarification regarding multiple authors
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  from typing import List
  from typing import Optional
  from typing import Tuple
  
  import torch
  import _k2
  import k2
  
  from .fsa import Fsa
  from .dense_fsa_vec import DenseFsaVec
  
  
  class _GetTotScoresFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx, fsas: Fsa, log_semiring: bool, use_double_scores: bool,
                  unused_scores: torch.Tensor) -> torch.Tensor:
          '''Compute the total loglikes of an FsaVec.
  
          Args:
            fsas:
              The input FsaVec (must have 3 axes)
            log_semiring:
              True to use log semiring, false to use tropical
            use_double_scores:
              False to use float, i.e., single precision floating point,
              to compute log likes. True to use double precision.
            unused_scores:
              It is used only for backward propagation purpose.
              It should equal `fsas.scores`.
  
          Returns:
             The total loglike for each FSA in `fsas`.  If
            `use_double_scores==True`, its dtype is `torch.float64`; it is
            `torch.float32` otherwise.
  
          '''
          # This function is called by fsas.get_tot_scores() and calls
          # fsas._get_tot_scores() (which is not differentiable).  the .detach()
          # below avoids a reference cycle; if we didn't do that, the backward_fn
          # of tot_scores, which is cached in `fsas`, would be set to this object,
          # giving `fsas` a reference to this object, which also has a reference
          # to `fsas`.
          tot_scores = fsas._get_tot_scores(use_double_scores=use_double_scores,
                                            log_semiring=log_semiring).detach()
  
          # NOTE: since `fsas`, `log_semiring` and `use_double_scores` are
          # not tensors, they are saved as attributes of `ctx`.
          ctx.fsas = fsas
          ctx.log_semiring = log_semiring
          ctx.use_double_scores = use_double_scores
  
          ctx.save_for_backward(unused_scores)
  
          return tot_scores
  
      @staticmethod
      def backward(ctx, tot_scores_grad: torch.Tensor
                  ) -> Tuple[None, None, None, torch.Tensor]:  # noqa
          """
          Caution: this backward function uses a slightly indirect approach to
          compute the gradients.  Since the tot_scores are just computed as
          specific elements of `forward_scores`, the obvious way to get
          derivatives w.r.t. fsas.scores would be to set gradients w.r.t. the
          forward scores and then use BackpropGetForwardScores() to do the
          backprop.  But that might be a little slower than what we actually do.
          What we actually do is to compute the backward scores and use them and
          the forward scores to compute the posteriors, and let the derivs be the
          (posterior in FSA * loss_deriv w.r.t. that FSA's tot_prob).  The
          result is the same, and the underlying C++ code is simpler.
          (BackpropGetForwardScores() was added in order to compute slightly
          more difficult objective functions, that depend on the individual
          arc posteriors).
          """
          fsas = ctx.fsas
          log_semiring = ctx.log_semiring
          use_double_scores = ctx.use_double_scores
          scores, = ctx.saved_tensors
  
          if log_semiring is False:
              entering_arcs = fsas._get_entering_arcs(use_double_scores)
              _, ragged_int = _k2.shortest_path(fsas.arcs, entering_arcs)
              if use_double_scores:
                  scores_grad = _k2.get_tot_scores_double_tropical_backward(
                      fsas.arcs, ragged_int, tot_scores_grad)
              else:
                  scores_grad = _k2.get_tot_scores_float_tropical_backward(
                      fsas.arcs, ragged_int, tot_scores_grad)
              # We return four values since the `forward` method accepts four
              # arguments (excluding ctx).
              #      fsas, log_semiring, use_double_scores, unused_scores
              return None, None, None, scores_grad
          else:
              arc_post = fsas._get_arc_post(use_double_scores, log_semiring)
              if use_double_scores:
                  bprop_func = _k2.get_tot_scores_double_log_backward
              else:
                  bprop_func = _k2.get_tot_scores_float_log_backward
              scores_grad = bprop_func(fsas.arcs, arc_post, tot_scores_grad)
              return None, None, None, scores_grad
  
  
  class _GetForwardScoresFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx, fsas: Fsa, log_semiring: bool, use_double_scores: bool,
                  unused_scores: torch.Tensor) -> torch.Tensor:
          '''Compute the forward scores of an FsaVec.
  
          Args:
            fsas:
              The input FsaVec (must have 3 axes)
            log_semiring:
              True to use log semiring, false to use tropical
            use_double_scores:
              False to use float, i.e., single precision floating point,
              to compute log likes. True to use double precision.
            unused_scores:
              It is used only for backward propagation purpose.
              It should equal `fsas.scores`.
  
          Returns:
             The total loglike for each FSA in `fsas`.  If
            `use_double_scores==True`, its dtype is `torch.float64`; it is
            `torch.float32` otherwise.
  
          '''
          # This function is called by fsas.get_forward_scores() and calls
          # fsas._get_forward_scores() (which is not differentiable).  the
          # .detach() below avoids a reference cycle, I believe; if we didn't do
          # that, the backward_fn of forward_scores, which is cached in `fsas`,
          # would be set to this object, giving `fsas` a reference to this object,
          # which also has a reference to `fsas`.
          forward_scores = fsas._get_forward_scores(
              use_double_scores=use_double_scores,
              log_semiring=log_semiring).detach()
  
          # NOTE: since `fsas`, `log_semiring` and `use_double_scores` are
          # not tensors, they are saved as attributes of `ctx`.
          ctx.fsas = fsas
          ctx.log_semiring = log_semiring
          ctx.use_double_scores = use_double_scores
          ctx.save_for_backward(forward_scores)
  
          return forward_scores
  
      @staticmethod
      def backward(ctx, forward_scores_grad: torch.Tensor
                  ) -> Tuple[None, None, None, torch.Tensor]:  # noqa
          fsas = ctx.fsas
          log_semiring = ctx.log_semiring
          use_double_scores = ctx.use_double_scores
          forward_scores, = ctx.saved_tensors
  
          if log_semiring:
              entering_arcs = None
          else:
              entering_arcs = fsas._get_entering_arcs(use_double_scores)
          state_batches = fsas._get_state_batches()
          leaving_arc_batches = fsas._get_leaving_arc_batches()
  
          bprop_func = (_k2.backprop_get_forward_scores_double
                        if use_double_scores else
                        _k2.backprop_get_forward_scores_float)
  
          scores_grad = bprop_func(fsas.arcs,
                                   state_batches=state_batches,
                                   leaving_arc_batches=leaving_arc_batches,
                                   log_semiring=log_semiring,
                                   entering_arcs=entering_arcs,
                                   forward_scores=forward_scores,
                                   forward_scores_deriv=forward_scores_grad)
  
          return (
              None,  # fsas
              None,  # log_semiring
              None,  # use_double_scores
              scores_grad  # unused_scores
          )
  
  
  class _GetBackwardScoresFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx, fsas: Fsa, log_semiring: bool, use_double_scores: bool,
                  unused_scores: torch.Tensor) -> torch.Tensor:
          '''Compute the backward scores of an FsaVec.
  
          Args:
            fsas:
              The input FsaVec (must have 3 axes)
            log_semiring:
              True to use log semiring, false to use tropical
            use_double_scores:
              False to use float, i.e., single precision floating point,
              to compute log likes. True to use double precision.
            unused_scores:
              It is used only for backward propagation purpose.
              It should equal `fsas.scores`.
  
          Returns:
            A torch.Tensor with shape equal to (num_states,)
          '''
          # This function is called by fsas.get_backward_scores() and calls
          # fsas._get_backward_scores() (which is not differentiable).  the
          # .detach() below avoids a reference cycle, I believe; if we didn't do
          # that, the backward_fn of backward_scores, which is cached in `fsas`,
          # would be set to this object, giving `fsas` a reference to this object,
          # which also has a reference to `fsas`.
          backward_scores = fsas._get_backward_scores(
              use_double_scores=use_double_scores,
              log_semiring=log_semiring).detach()
  
          # NOTE: since `fsas`, `log_semiring` and `use_double_scores` are
          # not tensors, they are saved as attributes of `ctx`.
          ctx.fsas = fsas
          ctx.log_semiring = log_semiring
          ctx.use_double_scores = use_double_scores
          ctx.save_for_backward(backward_scores)
  
          return backward_scores
  
      @staticmethod
      def backward(ctx, backward_scores_grad: torch.Tensor
                  ) -> Tuple[None, None, None, torch.Tensor]:  # noqa
          fsas = ctx.fsas
          log_semiring = ctx.log_semiring
          use_double_scores = ctx.use_double_scores
          backward_scores, = ctx.saved_tensors
  
          state_batches = fsas._get_state_batches()
          entering_arc_batches = fsas._get_entering_arc_batches()
  
          bprop_func = (_k2.backprop_get_backward_scores_double
                        if use_double_scores else
                        _k2.backprop_get_backward_scores_float)
  
          scores_grad = bprop_func(fsas.arcs,
                                   state_batches=state_batches,
                                   entering_arc_batches=entering_arc_batches,
                                   log_semiring=log_semiring,
                                   backward_scores=backward_scores,
                                   backward_scores_deriv=backward_scores_grad)
  
          return (
              None,  # fsas
              None,  # log_semiring
              None,  # use_double_scores
              scores_grad  # unused_scores
          )
  
  
  class _GetArcPostFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx, fsas: Fsa, log_semiring: bool, use_double_scores: bool,
                  unused_scores: torch.Tensor, forward_scores: torch.Tensor,
                  backward_scores: torch.Tensor) -> torch.Tensor:
          '''Compute the arc-level posteriors of an FsaVec
  
          Args:
            fsas:
              The input FsaVec (must have 3 axes)
            log_semiring:
              True to use log semiring, false to use tropical
            use_double_scores:
              False to use float, i.e., single precision floating point,
              to compute log likes. True to use double precision.
            unused_scores:
              It is used only for backward propagation purpose.
              It should equal `fsas.scores`.
            forward_scores:
              The forward scores of the FSA, computed in a differentiable
              way by fsas.get_forward_scores(); must be provided as an
              explicit arg for backprop reasons.
            backward_scores:
              The backward scores of the FSA, computed in a differentiable
              way from fsas.get_backward_scores(); must be provided as an
              explicit arg for backprop reasons.
  
          Returns:
            The per-arc log-posterior for each arc in `fsas`.  If
            `use_double_scores==True`, its dtype is `torch.float64`; it is
            `torch.float32` otherwise.
  
          '''
          # This function is called by fsas.get_arc_post() and calls
          # fsas._get_arc_post() (which is not differentiable) for caching
          # reasons, so the output can be cached there (although the backprop may
          # have to be repeated).  The .detach() below avoids a reference cycle;
          # if we didn't do that, the backward_fn of the arc_post, which is cached
          # in `fsas`, would be set to this object, giving `fsas` a reference to
          # this object, which also has a reference to `fsas`.
          arc_post = fsas._get_arc_post(use_double_scores=use_double_scores,
                                        log_semiring=log_semiring).detach()
  
          # NOTE: since `fsas`, `log_semiring` and `use_double_scores` are
          # not tensors, they are saved as attributes of `ctx`.
          ctx.fsas = fsas
          ctx.use_double_scores = use_double_scores
  
          ctx.save_for_backward(forward_scores, backward_scores)
          return arc_post
  
      @staticmethod
      def backward(
              ctx, arc_post_grad: torch.Tensor
      ) -> Tuple[None, None, None, torch.Tensor, torch.Tensor,  # noqa
                 torch.Tensor]:
          fsas = ctx.fsas
          use_double_scores = ctx.use_double_scores
          forward_scores, backward_scores = ctx.saved_tensors
  
          bprop_func = (_k2.backprop_get_arc_post_double if use_double_scores
                        else _k2.backprop_get_arc_post_float)
  
          incoming_arcs = fsas._get_incoming_arcs()
          forward_scores_grad, backward_scores_grad = bprop_func(
              fsas.arcs, incoming_arcs, arc_post_grad)
          arc_scores_grad = arc_post_grad.detach().clone()
  
          return (
              None,  # fsas
              None,  # log_semiring
              None,  # use_double_scores
              arc_scores_grad,  # unused_scores
              forward_scores_grad,  # forward_scores
              backward_scores_grad  # backward_scores
          )
  
  
  class _IntersectDensePrunedFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx,
                  a_fsas: Fsa,
                  b_fsas: DenseFsaVec,
                  out_fsa: List[Fsa],
                  search_beam: float,
                  output_beam: float,
                  min_active_states: int,
                  max_active_states: int,
                  unused_scores_a: torch.Tensor,
                  unused_scores_b: torch.Tensor,
                  seqframe_idx_name: Optional[str] = None,
                  frame_idx_name: Optional[str] = None) -> torch.Tensor:
          '''Intersect array of FSAs on CPU/GPU.
  
          Args:
            a_fsas:
              Input FsaVec, i.e., `decoding graphs`, one per sequence. It might
              just be a linear sequence of phones, or might be something more
              complicated. Must have either `a_fsas.shape[0] == b_fsas.dim0()`, or
              `a_fsas.shape[0] == 1` in which case the graph is shared.
            b_fsas:
              Input FSAs that correspond to neural network output.
            out_fsa:
              A list containing ONLY one entry which will be set to the
              generated FSA on return. We pass it as a list since the return
              value can only be types of torch.Tensor in the `forward` function.
            search_beam:
              Decoding beam, e.g. 20.  Smaller is faster, larger is more exact
              (less pruning). This is the default value; it may be modified by
              `min_active_states` and `max_active_states`.
            output_beam:
              Pruning beam for the output of intersection (vs. best path);
              equivalent to kaldi's lattice-beam.  E.g. 8.
            max_active_states:
              Maximum number of FSA states that are allowed to be active on any
              given frame for any given intersection/composition task. This is
              advisory, in that it will try not to exceed that but may not always
              succeed. You can use a very large number if no constraint is needed.
            min_active_states:
              Minimum number of FSA states that are allowed to be active on any
              given frame for any given intersection/composition task. This is
              advisory, in that it will try not to have fewer than this number
              active. Set it to zero if there is no constraint.
            unused_scores_a:
              It equals to `a_fsas.scores` and its sole purpose is for back
              propagation.
            unused_scores_b:
              It equals to `b_fsas.scores` and its sole purpose is for back
              propagation.
            seqframe_idx_name:
              If set (e.g. to 'seqframe'), an attribute in the output will be
              created that encodes the sequence-index and the frame-index within
              that sequence; this is equivalent to a row-index into b_fsas.values,
              or, equivalently, an element in b_fsas.shape.
            frame_idx_name:
              If set (e.g. to 'frame', an attribute in the output will be created
              that contains the frame-index within the corresponding sequence.
          Returns:
             Return `out_fsa[0].scores`.
          '''
          assert len(out_fsa) == 1
  
          ragged_arc, arc_map_a, arc_map_b = _k2.intersect_dense_pruned(
              a_fsas=a_fsas.arcs,
              b_fsas=b_fsas.dense_fsa_vec,
              search_beam=search_beam,
              output_beam=output_beam,
              min_active_states=min_active_states,
              max_active_states=max_active_states)
  
          out_fsa[0] = Fsa(ragged_arc)
  
          for name, a_value in a_fsas.named_tensor_attr(include_scores=False):
              if isinstance(a_value, torch.Tensor):
                  value = _k2.index_select(a_value, arc_map_a)
              else:
                  assert isinstance(a_value, k2.RaggedTensor)
                  # Only integer types ragged attributes are supported now
                  assert a_value.dtype == torch.int32
                  value, _ = a_value.index(arc_map_a,
                                           axis=0,
                                           need_value_indexes=False)
  
              setattr(out_fsa[0], name, value)
  
          for name, a_value in a_fsas.named_non_tensor_attr():
              setattr(out_fsa[0], name, a_value)
  
          ctx.arc_map_a = arc_map_a
          ctx.arc_map_b = arc_map_b
  
          ctx.save_for_backward(unused_scores_a, unused_scores_b)
  
          seqframe_idx = None
          if frame_idx_name is not None:
              num_cols = b_fsas.dense_fsa_vec.scores_dim1()
              seqframe_idx = arc_map_b // num_cols
              shape = b_fsas.dense_fsa_vec.shape()
              fsa_idx0 = _k2.index_select(shape.row_ids(1), seqframe_idx)
              frame_idx = seqframe_idx - _k2.index_select(
                  shape.row_splits(1), fsa_idx0)
              assert not hasattr(out_fsa[0], frame_idx_name)
              setattr(out_fsa[0], frame_idx_name, frame_idx)
  
          if seqframe_idx_name is not None:
              if seqframe_idx is None:
                  num_cols = b_fsas.dense_fsa_vec.scores_dim1()
                  seqframe_idx = arc_map_b // num_cols
  
              assert not hasattr(out_fsa[0], seqframe_idx_name)
              setattr(out_fsa[0], seqframe_idx_name, seqframe_idx)
  
          return out_fsa[0].scores
  
      @staticmethod
      def backward(ctx, out_fsa_grad: torch.Tensor) \
              -> Tuple[None, None, None, None, None, None, None, torch.Tensor, torch.Tensor]: # noqa
          a_scores, b_scores = ctx.saved_tensors
          arc_map_a = ctx.arc_map_a
          arc_map_b = ctx.arc_map_b
  
          grad_a = torch.zeros(a_scores.size(0),
                               dtype=out_fsa_grad.dtype,
                               device=a_scores.device,
                               requires_grad=False)
  
          grad_b = torch.zeros(
              *b_scores.shape,
              dtype=out_fsa_grad.dtype,
              device=b_scores.device,
              requires_grad=False).contiguous()  # will use its `view()` later
  
          _k2.index_add(arc_map_a, out_fsa_grad, grad_a)
          _k2.index_add(arc_map_b, out_fsa_grad, grad_b.view(-1))
  
          return (
              None,  # a_fass
              None,  # b_fsas
              None,  # out_fsa
              None,  # search_beam
              None,  # output_beam
              None,  # min_active_states
              None,  # max_active_states
              grad_a,  # unused_scores_a
              grad_b,  # unused_scores_b
              None,  # seqframe_idx_name
              None  # frame_idx_name
          )
  
  
  class _IntersectDenseFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx,
                  a_fsas: Fsa,
                  b_fsas: DenseFsaVec,
                  out_fsa: List[Fsa],
                  output_beam: float,
                  max_states: int,
                  max_arcs: int,
                  unused_scores_a: torch.Tensor,
                  unused_scores_b: torch.Tensor,
                  a_to_b_map: Optional[torch.Tensor] = None,
                  seqframe_idx_name: Optional[str] = None,
                  frame_idx_name: Optional[str] = None) -> torch.Tensor:
          '''Intersect array of FSAs on CPU/GPU.
  
          Args:
            a_fsas:
              Input FsaVec, i.e., `decoding graphs`, one per sequence. It might
              just be a linear sequence of phones, or might be something more
              complicated. Must have number of FSAs equal to b_fsas.dim0(), if
              a_to_b_map not specified.
            b_fsas:
              Input FSAs that correspond to neural network output.
            out_fsa:
              A list containing ONLY one entry which will be set to the
              generated FSA on return. We pass it as a list since the return
              value can only be types of torch.Tensor in the `forward` function.
            output_beam:
              Pruning beam for the output of intersection (vs. best path);
              equivalent to kaldi's lattice-beam.  E.g. 8.
            unused_scores_a:
              It equals to `a_fsas.scores` and its sole purpose is for back
              propagation.
            unused_scores_b:
              It equals to `b_fsas.scores` and its sole purpose is for back
              propagation.
            a_to_b_map:
              Maps from FSA-index in a to FSA-index in b to use for it.
              If None, then we expect the number of FSAs in a_fsas to equal
              b_fsas.dim0().  If set, then it should be a Tensor with ndim=1
              and dtype=torch.int32, with a_to_b_map.shape[0] equal to the
              number of FSAs in a_fsas (i.e. a_fsas.shape[0] if
              len(a_fsas.shape) == 3, else 1); and elements
              `0 <= i < b_fsas.dim0()`.
            seqframe_idx_name:
              If set (e.g. to 'seqframe'), an attribute in the output will be
              created that encodes the sequence-index and the frame-index within
              that sequence; this is equivalent to a row-index into b_fsas.values,
              or, equivalently, an element in b_fsas.shape.
            frame_idx_name:
              If set (e.g. to 'frame', an attribute in the output will be created
              that contains the frame-index within the corresponding sequence.
          Returns:
             Return `out_fsa[0].scores`.
          '''
          assert len(out_fsa) == 1
  
          ragged_arc, arc_map_a, arc_map_b = _k2.intersect_dense(
              a_fsas=a_fsas.arcs,
              b_fsas=b_fsas.dense_fsa_vec,
              a_to_b_map=a_to_b_map,
              output_beam=output_beam,
              max_states=max_states,
              max_arcs=max_arcs)
  
          out_fsa[0] = Fsa(ragged_arc)
  
          for name, a_value in a_fsas.named_tensor_attr(include_scores=False):
              if isinstance(a_value, torch.Tensor):
                  value = _k2.index_select(a_value, arc_map_a)
              else:
                  assert isinstance(a_value, k2.RaggedTensor)
                  assert a_value.dtype == torch.int32
                  value, _ = a_value.index(arc_map_a,
                                           axis=0,
                                           need_value_indexes=False)
  
              setattr(out_fsa[0], name, value)
  
          for name, a_value in a_fsas.named_non_tensor_attr():
              setattr(out_fsa[0], name, a_value)
  
          ctx.arc_map_a = arc_map_a
          ctx.arc_map_b = arc_map_b
  
          ctx.save_for_backward(unused_scores_a, unused_scores_b)
  
          seqframe_idx = None
          if frame_idx_name is not None:
              num_cols = b_fsas.dense_fsa_vec.scores_dim1()
              seqframe_idx = arc_map_b // num_cols
              shape = b_fsas.dense_fsa_vec.shape()
              fsa_idx0 = _k2.index_select(shape.row_ids(1), seqframe_idx)
              frame_idx = seqframe_idx - _k2.index_select(
                  shape.row_splits(1), fsa_idx0)
              assert not hasattr(out_fsa[0], frame_idx_name)
              setattr(out_fsa[0], frame_idx_name, frame_idx)
  
          if seqframe_idx_name is not None:
              if seqframe_idx is None:
                  num_cols = b_fsas.dense_fsa_vec.scores_dim1()
                  seqframe_idx = arc_map_b // num_cols
  
              assert not hasattr(out_fsa[0], seqframe_idx_name)
              setattr(out_fsa[0], seqframe_idx_name, seqframe_idx)
  
          return out_fsa[0].scores
  
      @staticmethod
      def backward(ctx, out_fsa_grad: torch.Tensor) \
              -> Tuple[None, None, None, None, torch.Tensor, torch.Tensor]: # noqa
          a_scores, b_scores = ctx.saved_tensors
          arc_map_a = ctx.arc_map_a
          arc_map_b = ctx.arc_map_b
  
          grad_a = torch.zeros(a_scores.size(0),
                               dtype=torch.float32,
                               device=a_scores.device,
                               requires_grad=False)
  
          grad_b = torch.zeros(
              *b_scores.shape,
              dtype=torch.float32,
              device=b_scores.device,
              requires_grad=False).contiguous()  # will use its `view()` later
  
          _k2.index_add(arc_map_a, out_fsa_grad, grad_a)
          _k2.index_add(arc_map_b, out_fsa_grad, grad_b.view(-1))
  
          return (
              None,  # a_fsas
              None,  # b_fsas
              None,  # out_fsa
              None,  # output_beam
              None,  # max_states
              None,  # max_arcs
              grad_a,  # unused_scores_a
              grad_b,  # unused_scores_b
              None,  # a_to_b_map
              None,  # seqframe_idx_name
              None  # frame_idx_name
          )
  
  
  def intersect_dense_pruned(a_fsas: Fsa,
                             b_fsas: DenseFsaVec,
                             search_beam: float,
                             output_beam: float,
                             min_active_states: int,
                             max_active_states: int,
                             seqframe_idx_name: Optional[str] = None,
                             frame_idx_name: Optional[str] = None) -> Fsa:
      '''Intersect array of FSAs on CPU/GPU.
  
      Caution:
        `a_fsas` MUST be arc sorted.
  
      Args:
        a_fsas:
          Input FsaVec, i.e., `decoding graphs`, one per sequence. It might just
          be a linear sequence of phones, or might be something more complicated.
          Must have either `a_fsas.shape[0] == b_fsas.dim0()`, or
          `a_fsas.shape[0] == 1` in which case the graph is shared.
        b_fsas:
          Input FSAs that correspond to neural network output.
        search_beam:
          Decoding beam, e.g. 20.  Smaller is faster, larger is more exact
          (less pruning). This is the default value; it may be modified by
          `min_active_states` and `max_active_states`.
        output_beam:
           Beam to prune output, similar to lattice-beam in Kaldi.  Relative
           to best path of output.
        min_active_states:
          Minimum number of FSA states that are allowed to be active on any given
          frame for any given intersection/composition task. This is advisory,
          in that it will try not to have fewer than this number active.
          Set it to zero if there is no constraint.
        max_active_states:
          Maximum number of FSA states that are allowed to be active on any given
          frame for any given intersection/composition task. This is advisory,
          in that it will try not to exceed that but may not always succeed.
          You can use a very large number if no constraint is needed.
        seqframe_idx_name:
          If set (e.g. to 'seqframe'), an attribute in the output will be created
          that encodes the sequence-index and the frame-index within that
          sequence; this is equivalent to a row-index into b_fsas.values,
          or, equivalently, an element in b_fsas.shape.
        frame_idx_name:
          If set (e.g. to 'frame', an attribute in the output will be created
          that contains the frame-index within the corresponding sequence.
  
      Returns:
        The result of the intersection.
      '''
      # Possible values for _k2.build_type are [Release, Debug]
      if _k2.version.build_type == 'Debug':
          # This check is to guarantee that all labels are in a valid range.
          # If not, unpredictable errors will occur.
          #
          # One such situation is that someone imports a graph from Kaldi,
          # whose labels are transition IDs. When the neural network output
          # units are pdf IDs, this additional check will detect the mismatch.
          #
          assert a_fsas.labels.min() >= -1
          # The first column of b_fsas.scores is -inf,
          # so we use b_fsas.scores.shape[1] - 1 here
          # (-1 is to exclude the column with -inf)
          assert a_fsas.labels.max() < b_fsas.scores.shape[1] - 1
  
      out_fsa = [0]
  
      # the following return value is discarded since it is already contained
      # in `out_fsa[0].scores`
      _IntersectDensePrunedFunction.apply(a_fsas, b_fsas, out_fsa, search_beam,
                                          output_beam, min_active_states,
                                          max_active_states, a_fsas.scores,
                                          b_fsas.scores, seqframe_idx_name,
                                          frame_idx_name)
      return out_fsa[0]
  
  
  def intersect_dense(a_fsas: Fsa,
                      b_fsas: DenseFsaVec,
                      output_beam: float,
                      max_states: int = 15000000,
                      max_arcs: int = 1073741824,
                      a_to_b_map: Optional[torch.Tensor] = None,
                      seqframe_idx_name: Optional[str] = None,
                      frame_idx_name: Optional[str] = None) -> Fsa:
      '''Intersect array of FSAs on CPU/GPU.
  
      Caution:
        `a_fsas` MUST be arc sorted.
  
      Args:
        a_fsas:
          Input FsaVec, i.e., `decoding graphs`, one per sequence. It might just
          be a linear sequence of phones, or might be something more complicated.
          Must have `a_fsas.shape[0] == b_fsas.dim0()` if `a_to_b_map` is None.
          Otherwise, must have `a_fsas.shape[0] == a_to_b_map.shape[0]`
        b_fsas:
          Input FSAs that correspond to neural network output.
        output_beam:
          Beam to prune output, similar to lattice-beam in Kaldi.  Relative
          to best path of output.
        max_states:
          The max number of states to prune the output, mainly to avoid
          out-of-memory and numerical overflow, default 15,000,000.
        max_arcs:
          The max number of arcs to prune the output, mainly to avoid
          out-of-memory and numerical overflow, default 1073741824(2^30).
        a_to_b_map:
           Maps from FSA-index in a to FSA-index in b to use for it.
           If None, then we expect the number of FSAs in a_fsas to equal
           b_fsas.dim0().  If set, then it should be a Tensor with ndim=1
           and dtype=torch.int32, with a_to_b_map.shape[0] equal to the
           number of FSAs in a_fsas (i.e. a_fsas.shape[0] if
           len(a_fsas.shape) == 3, else 1); and elements 0 <= i < b_fsas.dim0().
        seqframe_idx_name:
          If set (e.g. to 'seqframe'), an attribute in the output will be created
          that encodes the sequence-index and the frame-index within that
          sequence; this is equivalent to a row-index into b_fsas.values,
          or, equivalently, an element in b_fsas.shape.
        frame_idx_name:
          If set (e.g. to 'frame', an attribute in the output will be created
          that contains the frame-index within the corresponding sequence.
  
      Returns:
        The result of the intersection (pruned to `output_beam`; this pruning
        is exact, it uses forward and backward scores.
      '''
      # Possible values for _k2.build_type are [Release, Debug]
      if _k2.version.build_type == 'Debug':
          # This check is to guarantee that all labels are in a valid range.
          # If not, unpredictable errors will occur.
          #
          # One such situation is that someone imports a graph from Kaldi,
          # whose labels are transition IDs. When the neural network output
          # units are pdf IDs, this additional check will detect the mismatch.
          #
          assert a_fsas.labels.min() >= -1
          # The first column of b_fsas.scores is -inf,
          # so we use b_fsas.scores.shape[1] - 1 here
          # (-1 is to exclude the column with -inf)
          assert a_fsas.labels.max() < b_fsas.scores.shape[1] - 1
  
      out_fsa = [0]
  
      # the following return value is discarded since it is already contained
      # in `out_fsa[0].scores`
      _IntersectDenseFunction.apply(a_fsas, b_fsas, out_fsa, output_beam,
                                    max_states, max_arcs,
                                    a_fsas.scores, b_fsas.scores, a_to_b_map,
                                    seqframe_idx_name, frame_idx_name)
      return out_fsa[0]
  