o GicI@sddlZddlZddlZddlmZddlmZmZmZmZm Z Gdddej j Z  ddeded eed e d eeeeeeffeff d d Zdeded efddZ ddeedeed eed eefddZdS)N)Tensor)TupleOptionalSequenceUnionListc@szeZdZdZe  ddejdejdeeejdeejde d ejf d d Z ed ed e ejejdddffd dZ dS)"MutualInformationRecursionFunctionzA recursion that is useful in computing mutual information between two sequences of real vectors, but may be useful more generally in sequence-to-sequence tasks where monotonic alignment between pairs of sequences is desired. NFpxpy pxy_gradsboundary return_gradreturncCs$|j\}}}|jd} || | dfvsJ|| f|j||d| fks+J|j|| f|dur=|j|dfks=J|j|ftj||d| d|j|jd} t|||| } d\} } |sa|jsa|jr|tj||j|jd}t |||| |\} } | | | t |dksJt || |d<| |d<| S) a Computing mutual information between two sequences of real vectors. Args: px: A torch.Tensor of some floating point type, with shape ``[B][S][T]`` if modified, ``[B][S][T+1]`` if not modified. where ``B`` is the batch size, ``S`` is the length of the ``x`` sequence (including representations of ``EOS`` symbols but not ``BOS`` symbols), and ``T`` is the length of the ``y`` sequence (including representations of ``EOS`` symbols but not ``BOS`` symbols). In the mutual information application, ``px[b][s][t]`` would represent the following log odds ratio; ignoring the b index on the right to make the notation more compact:: px[b][s][t] = log [ p(x_s | x_{0..s-1}, y_{0..t-1}) / p(x_s) ] This expression also implicitly includes the log-probability of choosing to generate an ``x`` value as opposed to a ``y`` value. In practice it might be computed as ``a + b``, where ``a`` is the log probability of choosing to extend the sequence of length ``(s,t)`` with an ``x`` as opposed to a ``y`` value; and ``b`` might in practice be of the form:: log(N exp f(x_s, y_{t-1}) / sum_t' exp f(x_s, y_t')) where ``N`` is the number of terms that the sum over ``t'`` included, which might include some or all of the other sequences as well as this one. Note: we don't require ``px`` and ``py`` to be contiguous, but the code assumes for optimization purposes that the ``T`` axis has stride 1. py: A torch.Tensor of the same dtype as ``px``, with shape ``[B][S+1][T]``, representing:: py[b][s][t] = log [ p(y_t | x_{0..s-1}, y_{0..t-1}) / p(y_t) ] This function does not treat ``x`` and ``y`` differently; the only difference is that for optimization purposes we assume the last axis (the ``t`` axis) has stride of 1; this is true if ``px`` and ``py`` are contiguous. pxy_grads: A List to store the return grads of ``px`` and ``py`` if return_grad == True. Remain unchanged if return_grad == False. See `this PR ` for more information about why we add this parameter. Note: the length of the list must be 2, where the first element represents the grads of ``px`` and the second one represents the grads of ``py``. boundary: If supplied, a torch.LongTensor of shape ``[B][4]``, where each row contains ``[s_begin, t_begin, s_end, t_end]``, with ``0 <= s_begin <= s_end <= S`` and ``0 <= t_begin <= t_end < T`` (this implies that empty sequences are allowed). If not supplied, the values ``[0, 0, S, T]`` will be assumed. These are the beginning and one-past-the-last positions in the ``x`` and ``y`` sequences respectively, and can be used if not all sequences are of the same length. return_grad: Whether to return grads of ``px`` and ``py``, this grad standing for the occupation probability is the output of the backward with a ``fake gradient`` the ``fake gradient`` is the same as the gradient you'd get if you did ``torch.autograd.grad((scores.sum()), [px, py])``. This is useful to implement the pruned version of rnnt loss. Returns: Returns a torch.Tensor of shape ``[B]``, containing the log of the mutual information between the b'th pair of sequences. This is defined by the following recursion on ``p[b,s,t]`` (where ``p`` is of shape ``[B,S+1,T+1]``), representing a mutual information between sub-sequences of lengths ``s`` and ``t``:: p[b,0,0] = 0.0 p[b,s,t] = log_add(p[b,s-1,t] + px[b,s-1,t], p[b,s,t-1] + py[b,s,t-1]) (if s > 0 or t > 0) where we handle edge cases by treating quantities with negative indexes as **-infinity**. The extension to cases where the boundaries are specified should be obvious; it just works on shorter sequences with offsets into ``px`` and ``py``. Ndevicedtype)NNr) shapetorchemptyrr_k2mutual_information_forward requires_gradonesmutual_information_backwardsave_for_backwardlen)ctxr r r r r BST1Tpanspx_gradpy_gradans_gradr*I/home/ubuntu/.local/lib/python3.10/site-packages/k2/mutual_information.pyforwards& j $   z*MutualInformationRecursionFunction.forwardr)cCs>|j\}}|j\}||dd}||}||}||dddfS)Nr) saved_tensorsrreshape)r r)r'r(r!r*r*r+backwards z+MutualInformationRecursionFunction.backwardNF) __name__ __module__ __qualname____doc__ staticmethodrrrrboolr,rr/r*r*r*r+rs2  rFr r r r rcCs|jdks J|j|j\}}}|jd}|jd||dfvs'J|j|f|j||d|fks:J|j|||f|j|jksHJ|j|jf|dur|jtjksWJ|j|j|dfkseJ|j|f|D]4\}} } } d|kr~| kr~|ksnJ|| |fd| kr| kr|ksnJ| | |fqi||}}ddg} t||| ||} | \}}|r| ||ffS| S)aA recursion that is useful in computing mutual information between two sequences of real vectors, but may be useful more generally in sequence-to-sequence tasks where monotonic alignment between pairs of sequences is desired. The definitions of the arguments are definitions that would be used when computing this type of mutual information, but you can also view them as arbitrary quantities and just make use of the formula computed by this function. Args: px: A torch.Tensor of some floating point type, with shape ``[B][S][T+1]``, where ``B`` is the batch size, ``S`` is the length of the ``x`` sequence (including representations of ``EOS`` symbols but not ``BOS`` symbols), and ``T`` is the length of the ``y`` sequence (including representations of ``EOS`` symbols but not ``BOS`` symbols). In the mutual information application, ``px[b][s][t]`` would represent the following log odds ratio; ignoring the b index on the right to make the notation more compact:: px[b][s][t] = log [ p(x_s | x_{0..s-1}, y_{0..t-1}) / p(x_s) ] This expression also implicitly includes the log-probability of choosing to generate an ``x`` value as opposed to a ``y`` value. In practice it might be computed as ``a + b``, where ``a`` is the log probability of choosing to extend the sequence of length ``(s,t)`` with an ``x`` as opposed to a ``y`` value; and ``b`` might in practice be of the form:: log(N exp f(x_s, y_{t-1}) / sum_t' exp f(x_s, y_t')) where ``N`` is the number of terms that the sum over ``t'`` included, which might include some or all of the other sequences as well as this one. Note: we don't require ``px`` and ``py`` to be contiguous, but the code assumes for optimization purposes that the ``T`` axis has stride 1. py: A torch.Tensor of the same dtype as ``px``, with shape ``[B][S+1][T]``, representing:: py[b][s][t] = log [ p(y_t | x_{0..s-1}, y_{0..t-1}) / p(y_t) ] This function does not treat ``x`` and ``y`` differently; the only difference is that for optimization purposes we assume the last axis (the ``t`` axis) has stride of 1; this is true if ``px`` and ``py`` are contiguous. boundary: If supplied, a torch.LongTensor of shape ``[B][4]``, where each row contains ``[s_begin, t_begin, s_end, t_end]``, with ``0 <= s_begin <= s_end <= S`` and ``0 <= t_begin <= t_end < T`` (this implies that empty sequences are allowed). If not supplied, the values ``[0, 0, S, T]`` will be assumed. These are the beginning and one-past-the-last positions in the ``x`` and ``y`` sequences respectively, and can be used if not all sequences are of the same length. return_grad: Whether to return grads of ``px`` and ``py``, this grad standing for the occupation probability is the output of the backward with a ``fake gradient`` the ``fake gradient`` is the same as the gradient you'd get if you did ``torch.autograd.grad((scores.sum()), [px, py])``. This is useful to implement the pruned version of rnnt loss. Returns: Returns a torch.Tensor of shape ``[B]``, containing the log of the mutual information between the b'th pair of sequences. This is defined by the following recursion on ``p[b,s,t]`` (where ``p`` is of shape ``[B,S+1,T+1]``), representing a mutual information between sub-sequences of lengths ``s`` and ``t``:: p[b,0,0] = 0.0 if !modified: p[b,s,t] = log_add(p[b,s-1,t] + px[b,s-1,t], p[b,s,t-1] + py[b,s,t-1]) if modified: p[b,s,t] = log_add(p[b,s-1,t-1] + px[b,s-1,t-1], p[b,s,t-1] + py[b,s,t-1]) where we handle edge cases by treating quantities with negative indexes as **-infinity**. The extension to cases where the boundaries are specified should be obvious; it just works on shorter sequences with offsets into ``px`` and ``py``. rrNrr) ndimrrrint64tolist contiguousrapply)r r r r r!r"r#r$s_begint_begins_endt_endr scoresr'r(r*r*r+mutual_information_recursions&]  $&.0 rBabcCsT|jd|jdksJ|j|jf|d}|d}t||}|ddS)z Does inner product on the last dimension, with expected broadcasting, i.e. equivalent to (a * b).sum(dim=-1) without creating a large temporary. r)r unsqueezermatmulsqueeze)rCrDcr*r*r+_inner_product-s $   rJcCst|}t||kr|dksJt||f|dj\}}}|djd}|||dfvs3J||f|dj||d|fksJJ|dj|||f|dj|djks`J|dj|djftj|dd}tj|dd} |jdd} | jdd} |dur|jtjksJ|j|j|dfksJ|j|f|D]4\} } }}d| kr|kr|ksnJ| ||fd| kr|kr|ksnJ| ||fq| | } } | j dksJ| j| j dksJ| jtj ||d|d| j | jd}t | | ||}tj|| j | jd}t | | |||\}}|d|d }|d|d }|||d }| ||d } |jt|jjd }| jt| jjd } t||}t|| }||}t||jdd}Wdn 1srwY|d|7<|S) aA recursion that is useful for modifications of RNN-T and similar loss functions, where the recursion probabilities have a number of terms and you want them reported separately. See mutual_information_recursion() for more documentation of the basic aspects of this. Args: px: a sequence of Tensors, each of the same shape [B][S][T+1] py: a sequence of Tensor, each of the same shape [B][S+1][T], the sequence must be the same length as px. boundary: optionally, a LongTensor of shape [B][4] containing rows [s_begin, t_begin, s_end, t_end], with 0 <= s_begin <= s_end <= S and 0 <= t_begin <= t_end < T, defaulting to [0, 0, S, T]. These are the beginning and one-past-the-last positions in the x and y sequences respectively, and can be used if not all sequences are of the same length. Returns: a Tensor of shape (len(px), B), whose sum over dim 0 is the total log-prob of the recursion mentioned below, per sequence. The first element of the sequence of length len(px) is "special", in that it has an offset term reflecting the difference between sum-of-log and log-of-sum; for more interpretable loss values, the "main" part of your loss function should be first. The recursion below applies if boundary == None, when it defaults to (0, 0, S, T); where px_sum, py_sum are the sums of the elements of px and py:: p = tensor of shape (B, S+1, T+1), containing -infinity p[b,0,0] = 0.0 # do the following in loop over s and t: p[b,s,t] = log_add(p[b,s-1,t] + px_sum[b,s-1,t], p[b,s,t-1] + py_sum[b,s,t-1]) (if s > 0 or t > 0) return b[:][S][T] This function lets you implement the above recursion efficiently, except that it gives you a breakdown of the contribution from all the elements of px and py separately. As noted above, the first element of the sequence is "special". rrr)dimNrr7rr)min)rrrrstacksumr9r:r;r8rrrrrrr.clampfinforLrJno_grad)r r r Nr!r"r#r$px_catpy_catpx_totpy_totr=r>r?r@r% tot_probsr)r'r(x_prodsy_prodsprodsoffsetr*r*r+"joint_mutual_information_recursion:sX0$.,  .0     r\r0)N)osrrrtypingrrrrrautogradFunctionrr6rBrJr\r*r*r*r+s> % t