ujhՊ,@ddlmZmZddlmZmZmZddlZddlm Z gdZ GddejjZ Gdd ejjZGd d eZGd d ejjeZGddejjZGddejjZGddejjZdedededededededededededed ed!ed"ed#ed$ed%ed&ed'ed(ed)ef,d*Zded)efd+Zy),)ABCabstractmethod)ListOptionalTupleN)Emformer)RNNTemformer_rnnt_baseemformer_rnnt_modelceZdZdZdeddffd ZdejdejdeejejffdZ xZ S) _TimeReductionzCoalesces frames along time dimension into a fewer number of frames with higher feature dimensionality. Args: stride (int): number of frames to merge for each output frame. stridereturnNc0t|||_yN)super__init__r)selfr __class__s T/var/www/html/dev/engine/venv/lib/python3.12/site-packages/torchaudio/models/rnnt.pyrz_TimeReduction.__init__s  inputlengthsc"|j\}}}|||jzz }|ddd|ddf}|j|jd}||jz}|j||||jz}|j }||fS)aForward pass. B: batch size; T: maximum input sequence length in batch; D: feature dimension of each input sequence frame. Args: input (torch.Tensor): input sequences, with shape `(B, T, D)`. lengths (torch.Tensor): with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in ``input``. Returns: (torch.Tensor, torch.Tensor): torch.Tensor output sequences, with shape `(B, T // stride, D * stride)` torch.Tensor output lengths, with shape `(B,)` and i-th element representing number of valid frames for i-th batch element in output sequences. Ntrunc) rounding_mode)shaperdivreshape contiguous) rrrBTD num_framesT_maxoutputs rforwardz_TimeReduction.forwards*++1a!dkk/* a*a'(++dkk+Adkk)q%T[[9""$wr) __name__ __module__ __qualname____doc__intrtorchTensorrr' __classcell__rs@rr r sRstU\\ELLU5<