Speeding up my logsumexp function

General Usage Performance
, , ,
2

Use maximum to find the maximum:

function lsexp_mat2(mat, zero1_mat; dims=1)

    # To use the function zero1_mat has to be created outside @model
    # with the following code
    # zero1_mat = zeros(size(input_matrix))
    # zero1_mat[end, :] = zero1_mat[end, :] .+ 1

    max_ = maximum(mat, dims=dims)
    exp_mat = exp.(mat .- max_)
    sum_exp_ = sum(exp_mat .- zero1_mat, dims=dims)
    log1p.(sum_exp_) .+ max_
end