The sum(f, x) pattern may be a bit of a waste, as IIRC Zygote turns it back into sum(f.(x)) (whereas normally it would save allocating f.(x)).
I think so? I mean I suggest working with “solid” arrays as much as possible, rather than making slices. So instead of (or in addition to) istft(x::AbstractMatrix) you define istft(xs::AbstractArray{T,3}) with every step inside this function keeping track of one more dimension.