I wrote a logsumexp function based on the code here, that takes matrix as input and applies logsumexp along a dimension in a numerically stable way.
function lsexp_mat(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 sorted_mat = sort(mat, dims=dims) max_ = reshape(sorted_mat[end, :], 1, :) exp_mat = exp.(sorted_mat .- max_) sum_exp_ = sum(exp_mat .- zero1_mat, dims=dims) log1p.(sum_exp_) .+ max_ end
sort function to make it easy to subtract
1 max elements along the given dimension. But sorting is nlogn. I also wrote this way because I don’t any variables in my code to mutate because I want to use this code as a part of my Turing model and then use Zygote backend.
Is there way to create a mask matrix with
1 at all argmax indices and
0 else where without mutation?
This is the function from the discussion I’m trying to implement for a matrix.
function logsumexp!(p::WeightedParticles) N = length(p) w = p.logweights offset, maxind = findmax(w) w .= exp.(w .- offset) Σ = sum_all_but(w,maxind) # Σ = ∑wₑ-1 log1p(Σ) + offset, Σ+1 end function sum_all_but(w,i) w[i] -= 1 s = sum(w) w[i] += 1 s end