text: # Copyright    2021  Xiaomi Corp.        (authors: Fangjun Kuang)
  #
  # 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.
  
  import k2
  import torch
  import torch.nn as nn
  from encoder_interface import EncoderInterface
  
  from icefall.utils import add_sos
  
  
  class Transducer(nn.Module):
      """It implements https://arxiv.org/pdf/1211.3711.pdf
      "Sequence Transduction with Recurrent Neural Networks"
      """
  
      def __init__(
          self,
          encoder: EncoderInterface,
          decoder: nn.Module,
          joiner: nn.Module,
      ):
          """
          Args:
            encoder:
              It is the transcription network in the paper. Its accepts
              two inputs: `x` of (N, T, C) and `x_lens` of shape (N,).
              It returns two tensors: `logits` of shape (N, T, C) and
              `logit_lens` of shape (N,).
            decoder:
              It is the prediction network in the paper. Its input shape
              is (N, U) and its output shape is (N, U, C). It should contain
              one attribute: `blank_id`.
            joiner:
              It has two inputs with shapes: (N, T, C) and (N, U, C). Its
              output shape is (N, T, U, C). Note that its output contains
              unnormalized probs, i.e., not processed by log-softmax.
          """
          super().__init__()
          assert isinstance(encoder, EncoderInterface), type(encoder)
          assert hasattr(decoder, "blank_id")
  
          self.encoder = encoder
          self.decoder = decoder
          self.joiner = joiner
  
      def forward(
          self,
          x: torch.Tensor,
          x_lens: torch.Tensor,
          y: k2.RaggedTensor,
      ) -> torch.Tensor:
          """
          Args:
            x:
              A 3-D tensor of shape (N, T, C).
            x_lens:
              A 1-D tensor of shape (N,). It contains the number of frames in `x`
              before padding.
            y:
              A ragged tensor with 2 axes [utt][label]. It contains labels of each
              utterance.
          Returns:
            Return the transducer loss.
          """
          assert x.ndim == 3, x.shape
          assert x_lens.ndim == 1, x_lens.shape
          assert y.num_axes == 2, y.num_axes
  
          assert x.size(0) == x_lens.size(0) == y.dim0
  
          encoder_out, x_lens = self.encoder(x, x_lens)
          assert torch.all(x_lens > 0)
  
          # Now for the decoder, i.e., the prediction network
          row_splits = y.shape.row_splits(1)
          y_lens = row_splits[1:] - row_splits[:-1]
  
          blank_id = self.decoder.blank_id
          sos_y = add_sos(y, sos_id=blank_id)
  
          sos_y_padded = sos_y.pad(mode="constant", padding_value=blank_id)
  
          decoder_out = self.decoder(sos_y_padded)
  
          logits = self.joiner(encoder_out, decoder_out)
  
          # rnnt_loss requires 0 padded targets
          # Note: y does not start with SOS
          y_padded = y.pad(mode="constant", padding_value=0)
  
          y_padded = y_padded.to(torch.int64)
          boundary = torch.zeros(
              (x.size(0), 4), dtype=torch.int64, device=x.device
          )
          boundary[:, 2] = y_lens
          boundary[:, 3] = x_lens
  
          loss = k2.rnnt_loss(logits, y_padded, blank_id, boundary)
  
          return loss
  