Pytorch tensor plays

pytorch_twisted.py

import torch


def cum_softmax(t: torch.Tensor) -> torch.Tensor:
    # t shape: ..., sm_d, sm_d is the dim to reduce
    tmax = t.cummax(dim=-1)[0].unsqueeze(-2)
    denominator = (t.unsqueeze(-1) - tmax).exp()
    # shape: ..., sm_d, csm_d
    numerator = denominator.sum(dim=-2, keepdim=True)
    return numerator / denominator