Skip to content

A lesson in numerical stability #484

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bvdmitri merged 5 commits into from
Jul 4, 2025
Merged

A lesson in numerical stability #484

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b5142b2
improve memory footprint of DT
wouterwln Jun 20, 2025
7f42eb1
Improve numerical stability of DT rules
wouterwln Jul 4, 2025
f37e12b
Fix tests
wouterwln Jul 4, 2025
3c2d74f
🤖 Auto-format Julia code
wouterwln Jul 4, 2025
62682b0
Merge pull request #485 from ReactiveBayes/auto-format-484
wouterwln Jul 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/rules/discrete_transition/categoricals.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@ Compute the message for one of the Categorical interfaces of the `DiscreteTransi
function discrete_transition_structured_message_rule(message_names, messages, marginals_names, marginals, q_a)
e_log_a = mean(BroadcastFunction(clamplog), q_a)
e_log_a = discrete_transition_process_marginals(e_log_a, marginals_names, marginals)
msg = clamp.(exp.(e_log_a), tiny, huge)
msg = clamp.(softmax!(e_log_a), tiny, huge)
msg = discrete_transition_process_messages(msg, message_names, messages, sum_out_dimensions)
msg = reshape(msg, :)
normalize!(msg, 1)
Expand Down
13 changes: 8 additions & 5 deletions src/rules/discrete_transition/marginals.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,13 @@
::Any,
::Any
) where {marginal_symbol, message_names, N}
return Contingency(outer_product(probvec.(messages)) .* exp.(mean(BroadcastFunction(clamplog), first(marginals))))
result = outer_product(probvec.(messages)) .* softmax!(mean(BroadcastFunction(clamplog), first(marginals)))
normalize!(result, 1)
return Contingency(result, Val(false))

Check warning on line 18 in src/rules/discrete_transition/marginals.jl

View check run for this annotation

Codecov / codecov/patch

src/rules/discrete_transition/marginals.jl#L16-L18 

Added lines #L16 - L18 were not covered by tests
end

nonparametric_distribution(v::Vector{<:Real}) = Categorical(normalize!(v, 1))
nonparametric_distribution(v::AbstractArray{<:Real, N} where {N}) = Contingency(v)
nonparametric_distribution(v::Vector{<:Real}) = Categorical(normalize!(v, 1); check_args = false)
nonparametric_distribution(v::AbstractArray{<:Real, N} where {N}) = Contingency(normalize!(v, 1), Val(false))

# Generic implementation
"""
Expand All @@ -39,15 +41,16 @@
e_log_a = mean(BroadcastFunction(clamplog), q_a)
e_log_a = discrete_transition_process_marginals(e_log_a, marginals_names, marginals)

marginal = clamp.(exp.(e_log_a), tiny, huge)
marginal = clamp.(softmax!(e_log_a), tiny, huge)
marginal = discrete_transition_process_messages(marginal, message_names, messages, multiply_dimensions!)
dims = Tuple(findall(size(marginal) .== 1))
marginal = dropdims(marginal, dims = dims)
normalize!(marginal, 1)
return marginal
end

discrete_transition_marginal_rule_contingency(message_names::NTuple{N, Symbol}, messages::NTuple{N, Union{<:Message{<:DiscreteNonParametric}, <:Message{<:Bernoulli}}}, marginals_names::NTuple{M, Symbol}, marginals, q_a) where {N, M} = Contingency(
discrete_transition_marginal_rule(message_names, messages, marginals_names, marginals, q_a)
discrete_transition_marginal_rule(message_names, messages, marginals_names, marginals, q_a), Val(false)
)

function marginalrule(
Expand Down
5 changes: 3 additions & 2 deletions src/rules/discrete_transition/predefined/a.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ end
@rule DiscreteTransition(:a, Marginalisation) (q_out_in::Contingency, q_T1::PointMass{<:AbstractVector{T}}, meta::Any) where {T} = begin
out_in = components(q_out_in)
T1 = probvec(q_T1)
@tullio result[a, b, c] := out_in[a, b] * T1[c]
return DirichletCollection(result .+ 1)
result = ones(T, size(out_in)..., length(T1))
result[:, :, findfirst(isone, T1)] .+= out_in
return DirichletCollection(result)
end
80 changes: 48 additions & 32 deletions src/rules/discrete_transition/predefined/belief_propagation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,45 +2,52 @@ using Tullio

# # --------------- Rules for 2 interfaces (PointMass q_a) ---------------
@rule DiscreteTransition(:out, Marginalisation) (m_in::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 2}}, meta::Any) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_in))
eloga = clamp.(mean(q_a), tiny(N), one(N))
out = eloga * probvec(m_in)

return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (m_out::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 2}}, meta::Any) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
out = eloga' * probvec(m_out)
return Categorical(normalize!(out, 1); check_args = false)
end

# --------------- Rules for 2 interfaces (DirichletCollection q_a) ---------------
@rule DiscreteTransition(:out, Marginalisation) (m_in::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
N = eltype(probvec(m_in))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
out = eloga * probvec(m_in)
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (m_out::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
out = eloga' * probvec(m_out)
return Categorical(normalize!(out, 1); check_args = false)
end

# --------------- Rules for 3 interfaces (PointMass q_a) ---------------
@rule DiscreteTransition(:out, Marginalisation) (m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 3}}, meta::Any) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_in))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[i, a, b] * probvec(m_in)[a] * probvec(m_T1)[b]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (m_out::DiscreteNonParametric, m_T1::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 3}}, meta::Any) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, i, b] * probvec(m_out)[a] * probvec(m_T1)[b]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T1, Marginalisation) (m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 3}}, meta::Any) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, b, i] * probvec(m_out)[a] * probvec(m_in)[b]
return Categorical(normalize!(out, 1); check_args = false)
end
Expand All @@ -49,34 +56,34 @@ end
@rule DiscreteTransition(:out, Marginalisation) (m_in::DiscreteNonParametric, q_a::DirichletCollection, q_T1::PointMass{<:AbstractArray{T, 3}}, meta::Any) where {T} = begin
eloga = mean(Base.Broadcast.BroadcastFunction(clamplog), q_a)
@tullio intermediate[i, a] := eloga[i, a, b] * probvec(q_T1)[b]
out .= exp.(intermediate)
result = out * probvec(m_in)
softmax!(intermediate)
result = intermediate * probvec(m_in)
return Categorical(normalize!(result, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (m_out::DiscreteNonParametric, q_a::DirichletCollection, q_T1::PointMass{<:AbstractArray{T, 3}}, meta::Any) where {T} = begin
eloga = mean(Base.Broadcast.BroadcastFunction(clamplog), q_a)
@tullio intermediate[a, i] := eloga[a, i, b] * probvec(q_T1)[b]
out .= exp.(intermediate)
result = out' * probvec(m_out)
softmax!(intermediate)
result = intermediate' * probvec(m_out)
return Categorical(normalize!(result, 1); check_args = false)
end

# --------------- Rules for 3 interfaces (DirichletCollection q_a) ---------------
@rule DiscreteTransition(:out, Marginalisation) (m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[i, a, b] * probvec(m_in)[a] * probvec(m_T1)[b]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (m_out::DiscreteNonParametric, m_T1::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, i, b] * probvec(m_out)[a] * probvec(m_T1)[b]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T1, Marginalisation) (m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, b, i] * probvec(m_out)[a] * probvec(m_in)[b]
return Categorical(normalize!(out, 1); check_args = false)
end
Expand All @@ -85,56 +92,60 @@ end
@rule DiscreteTransition(:out, Marginalisation) (
m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 4}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_in))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[i, a, b, c] * probvec(m_in)[a] * probvec(m_T1)[b] * probvec(m_T2)[c]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (
m_out::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 4}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, i, b, c] * probvec(m_out)[a] * probvec(m_T1)[b] * probvec(m_T2)[c]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T1, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 4}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, b, i, c] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T2)[c]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T2, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 4}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, b, c, i] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T1)[c]
return Categorical(normalize!(out, 1); check_args = false)
end

# --------------- Rules for 4 interfaces (DirichletCollection q_a) ---------------
@rule DiscreteTransition(:out, Marginalisation) (m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[i, a, b, c] * probvec(m_in)[a] * probvec(m_T1)[b] * probvec(m_T2)[c]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (m_out::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, i, b, c] * probvec(m_out)[a] * probvec(m_T1)[b] * probvec(m_T2)[c]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T1, Marginalisation) (m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, b, i, c] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T2)[c]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T2, Marginalisation) (m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, b, c, i] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T1)[c]
return Categorical(normalize!(out, 1); check_args = false)
end
Expand All @@ -143,39 +154,44 @@ end
@rule DiscreteTransition(:out, Marginalisation) (
m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, m_T3::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 5}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_in))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[i, a, b, c, d] * probvec(m_in)[a] * probvec(m_T1)[b] * probvec(m_T2)[c] * probvec(m_T3)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (
m_out::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, m_T3::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 5}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, i, b, c, d] * probvec(m_out)[a] * probvec(m_T1)[b] * probvec(m_T2)[c] * probvec(m_T3)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T1, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T2::DiscreteNonParametric, m_T3::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 5}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, b, i, c, d] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T2)[c] * probvec(m_T3)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T2, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 5}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, b, c, i, d] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T1)[c] * probvec(m_T2)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T3, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::PointMass{<:AbstractArray{T, 5}}, meta::Any
) where {T} = begin
eloga = mean(q_a)
N = eltype(probvec(m_out))
eloga = clamp.(mean(q_a), tiny(N), one(N))
@tullio out[i] := eloga[a, b, c, d, i] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T1)[c] * probvec(m_T2)[d]
return Categorical(normalize!(out, 1); check_args = false)
end
Expand All @@ -184,39 +200,39 @@ end
@rule DiscreteTransition(:out, Marginalisation) (
m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, m_T3::DiscreteNonParametric, q_a::DirichletCollection, meta::Any
) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[i, a, b, c, d] * probvec(m_in)[a] * probvec(m_T1)[b] * probvec(m_T2)[c] * probvec(m_T3)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:in, Marginalisation) (
m_out::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, m_T3::DiscreteNonParametric, q_a::DirichletCollection, meta::Any
) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, i, b, c, d] * probvec(m_out)[a] * probvec(m_T1)[b] * probvec(m_T2)[c] * probvec(m_T3)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T1, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T2::DiscreteNonParametric, m_T3::DiscreteNonParametric, q_a::DirichletCollection, meta::Any
) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, b, i, c, d] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T2)[c] * probvec(m_T3)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T2, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::DirichletCollection, meta::Any
) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio out[i] := eloga[a, b, c, i, d] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T1)[c] * probvec(m_T2)[d]
return Categorical(normalize!(out, 1); check_args = false)
end

@rule DiscreteTransition(:T3, Marginalisation) (
m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, m_T1::DiscreteNonParametric, m_T2::DiscreteNonParametric, q_a::DirichletCollection, meta::Any
) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a), dims = 1)
@tullio out[i] := eloga[a, b, c, d, i] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T1)[c] * probvec(m_T2)[d]
return Categorical(normalize!(out, 1); check_args = false)
end
17 changes: 10 additions & 7 deletions src/rules/discrete_transition/predefined/marginals.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,21 +1,24 @@
using Tullio

@marginalrule DiscreteTransition(:out_in) (m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio result[a, b] := eloga[a, b] * probvec(m_out)[a] * probvec(m_in)[b]
return Contingency(result)
normalize!(result, 1)
return Contingency(result, Val(false))
end

@marginalrule DiscreteTransition(:out_in) (m_out::DiscreteNonParametric, m_in::DiscreteNonParametric, q_a::DirichletCollection, q_T1::PointMass, meta::Any) = begin
eloga = mean(Base.Broadcast.BroadcastFunction(clamplog), q_a)
@tullio result[a, b] := eloga[a, b, i] * probvec(q_T1)[i]
result = exp.(result)
result = eloga[:, :, findfirst(isone, probvec(q_T1))]
softmax!(result)
@tullio result[a, b] = result[a, b] * probvec(m_out)[a] * probvec(m_in)[b]
return Contingency(result)
normalize!(result, 1)
return Contingency(result, Val(false))
end

@marginalrule DiscreteTransition(:out_in_T1) (m_out::Categorical, m_in::Categorical, m_T1::Categorical, q_a::DirichletCollection, meta::Any) = begin
eloga = exp.(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
eloga = softmax!(mean(Base.Broadcast.BroadcastFunction(clamplog), q_a))
@tullio result[a, b, c] := eloga[a, b, c] * probvec(m_out)[a] * probvec(m_in)[b] * probvec(m_T1)[c]
return Contingency(result)
normalize!(result, 1)
return Contingency(result, Val(false))
end
Loading