(ctx, doys: torch.Tensor, *args)
| 681 | @staticmethod |
| 682 | @torch.cuda.amp.custom_bwd |
| 683 | def backward(ctx, doys: torch.Tensor, *args): |
| 684 | mask, us, dts, As, Bs, Cs, Ds, delta_bias, ohs = ctx.saved_tensors |
| 685 | |
| 686 | B, G, L, N, D = ctx.shape |
| 687 | chunksize = ctx.chunksize |
| 688 | delta_softplus = ctx.delta_softplus |
| 689 | doys = doys.view(B, G, D, L).permute(0, 1, 3, 2) |
| 690 | def rev_comsum_dim_2(x): |
| 691 | cum_sum = torch.cumsum(x, dim=2) |
| 692 | return (x - cum_sum + cum_sum[:,:,-1:None]) |
| 693 | |
| 694 | dus = None |
| 695 | dDs = None |
| 696 | if Ds is not None: |
| 697 | dDs = torch.einsum("bgld,bgld->gd", doys, us).view(-1) |
| 698 | dus = torch.einsum("bgld,gd->bgld", doys, Ds) |
| 699 | |
| 700 | chunks = list(range(0, L, chunksize)) |
| 701 | dAs = us.new_zeros((G, N, D), dtype=torch.float) |
| 702 | dus = us.new_zeros((B, G, L, D), dtype=torch.float) if dus is None else dus |
| 703 | ddts = us.new_zeros((B, G, L, D), dtype=torch.float) |
| 704 | dBs = us.new_zeros((B, G, L, N), dtype=torch.float) |
| 705 | dCs = us.new_zeros((B, G, L, N), dtype=torch.float) |
| 706 | dhprefix = us.new_zeros((B, G, N, D), dtype=torch.float) |
| 707 | |
| 708 | ohs_ptr = -2 |
| 709 | for i in chunks[::-1]: |
| 710 | h0 = ohs[:,:, ohs_ptr] |
| 711 | ohs_ptr = ohs_ptr - 1 |
| 712 | # forward procedure ================ |
| 713 | # BHLD, BHLN, HND, BHDN |
| 714 | cus = us[:,:,i:i + chunksize] |
| 715 | cdts = dts[:,:,i:i + chunksize] |
| 716 | cBs = Bs[:,:,i:i + chunksize] |
| 717 | cCs = Cs[:,:,i:i + chunksize] |
| 718 | cdoys = doys[:,:,i:i + chunksize] |
| 719 | cL = cus.shape[2] |
| 720 | _mask = (mask[:cL,:cL].contiguous() if cL < chunksize else mask).view(1, 1, cL, cL, 1) |
| 721 | |
| 722 | ts = torch.cumsum(cdts, dim=2) |
| 723 | w_log = As[None, :, None, :, :] * (ts[..., None, :]) # (B, H, L, Dk, Dv) |
| 724 | v = cus * cdts # (B,H,L,Dv) |
| 725 | k = cBs # (B,H,L,Dk) |
| 726 | q = cCs # (B,H,L,Dk) |
| 727 | w = w_log.exp() |
| 728 | |
| 729 | k_div_w = k[..., None] / w |
| 730 | q_mul_w = q[..., None] * w |
| 731 | |
| 732 | # h0 independent ================ |
| 733 | next_h_1_tmp = torch.einsum("bhlkv,bhlv->bhkv", k_div_w, v) |
| 734 | # next_h_1 = w[:,:,-1] * next_h_1_tmp |
| 735 | y_1_tmp = torch.einsum("bhlkv,bhrkv->bhlrv", q_mul_w, k_div_w) * _mask |
| 736 | # y_1 = torch.einsum("bhlrv,bhrv->bhlv", y_1_tmp, v) |
| 737 | |
| 738 | # h0 dependent ================ |
| 739 | # next_h_0 = w[:,:, -1] * h0 |
| 740 | # y_0 = torch.einsum("bhlkv,bhkv->bhlv", q_mul_w, h0) |
nothing calls this directly
no outgoing calls
no test coverage detected