Best way to do a “shifted” matrix multiplication

[marcofrancis]

I have two arrays say f and g, f is N by T by J dimensional and f is T by J dimensional. I’m trying to compute the following in JAX (for all 0<=t<T):

Notice that if t-a<0 I’d like it to default to 0.
What would be the fastest approach?

Right now I create a list of all possible indexes, multiply elementwise the two arrays evaluated in the relevant indexes and sum them up:

import jax.numpy as jnp

all_indices = jnp.array([(θ, t, a) for θ in range(N) for t in range(T) for a in range(J)])
θ_idx, t_idx, a_idx = all_indices[:, 0], all_indices[:, 1], all_indices[:, 2]
tma_idx = jnp.maximum(t_idx - a_idx, 0)

unrolled = f[θ_idx, t_idx, a_idx] * g[tma_idx, a_idx]
s = unrolled.reshape(N, T, J).sum(axis=(0,2))

This does not seem particularly efficient nor elegant and I would appreciate a better solution.

[jakevdp]

I suspect the best way to accomplish this is to first shift the 2D matrix, then perform the full reduction via einsum. For example:

t = jnp.arange(T)[:, None]
a = jnp.arange(J)
g_shifted = g[jnp.maximum(t - a, 0), a]
s = jnp.einsum("nta,ta->t", f, g_shifted)

Compared to the original solution, this reduces the number of index operations into g by a factor of T, eliminates the need to index into f, and relies on an efficient einsum operation to compute the final result.