Speeding up my logsumexp function

[vkv]

@vkv

function lsexp2(mat; dims=1)
    max_ = vreduce(max, mat, dims=dims)
    @tullio exp_mat[i,j] := exp(mat[i,j] - max_[1,j]) - (mat[i,j] == max_[1,j])
    sum_exp_ = sum(exp_mat, dims=dims)
    @tullio out[i,j] := log1p(sum_exp_[i,j]) + max_[i,j]
end

This function is the fastest one, but threw this error with mwe .

Is the gensym problem related to this error too? I thought ReverseDiff doesn’t support vreduce. Is there anyway to speed up maximum