Vectorizing adstock/geometric-decay operations for speed-up

I am trying to make this operation below a bit faster, does anyone have a smart way of doing so utilizing vectorized operations?

Also it would be nice to know how to initialize a tensor that has size 0, that case should also work in the below setup(see ones and hundres variables, imagine another variable with the below setup):

def adstock_geometric(x: torch.Tensor, theta: float):
    x_decayed = torch.zeros_like(x)
    x_decayed[0] = x[0]

    for xi in range(1, len(x_decayed)):
        x_decayed[xi] = x[xi] + theta * x_decayed[xi - 1]

    return x_decayed

def adstock_multiple_samples(x: torch.Tensor, theta: torch.Tensor):

    listtheta = theta.tolist()
    if isinstance(listtheta, float):
        return adstock_geometric(x=x,
    x_decayed = torch.zeros((100, 112, 1))
    for idx, theta_ in enumerate(listtheta):
        x_decayed_one_entry = adstock_geometric(x=x,
        x_decayed[idx] = x_decayed_one_entry
    return x_decayed

if __name__ == '__main__':
      ones = torch.tensor([1])
      hundreds = torch.tensor([idx for idx in range(100)])
      x = torch.tensor([[idx] for idx in range(112)])
      ones = adstock_multiple_samples(x=x,
      hundreds = adstock_multiple_samples(x=x,

Hello Richard,

So to give a more detailed answer, we’d need to know something about theta and the tensor shapes involved, so that remains the big If to the discussion below:

It’s a bit touchy w.r.t. numerical stability, but I think if you precomputed x_scaled = x * (theta ** (-xi)), you could use x_decayed = x_scaled.cumsum(0) * (theta ** (xi)) (with xi = torch.arange(...). It really depends on the size if xi and the magnitude of theta if this has a stability problem. In that case, you could try to split the computation in several parts and combine the vectorized approach with a for loop similar to what you have now to balance stability and speed.

Best regards


Hi Tom,
Thanks for replying.
theta.size() == () # True
theta.size() == (100,) # True
x.size() == (sizex,) where sizex [1…1000]. # True
are the ones involved.

kind regards,


But what is a typical value of theta?

sampled from a beta distribution so in the range [0, 1]

Yeah, so for close-to-one values of beta, the above would likely work, not sure about small ones.