Improve type stability of expectation

JuliaAlgebra · Nov 28, 2023 · 6e2df03 · 6e2df03
1 parent 442c962
commit 6e2df03
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/src/expectation.jl b/src/expectation.jl
@@ -1,6 +1,6 @@
-function _expectation(μ::Measure, p::_APL, f)
+function _expectation(μ::Measure{S}, p::_APL{T}, f) where {S,T}
     i = 1
-    s = 0
+    s = zero(MA.promote_operation(*, S, T))
     for t in MP.terms(p)
         while i <= length(μ.x) && MP.monomial(t) != μ.x[i]
             i += 1

diff --git a/test/expectation.jl b/test/expectation.jl
@@ -3,10 +3,10 @@
     p = x[3] - 2x[1] * x[2]^2 + 3x[3] * x[1] - 5x[1]^3
     v = (1, 2, 3)
     m = dirac(monomials(p), x => v)
-    @test MultivariateMoments.expectation(m, p) ==
+    @test (@inferred MultivariateMoments.expectation(m, p)) ==
           p(x => v) ==
-          MultivariateMoments.expectation(p, m)
+          (@inferred MultivariateMoments.expectation(p, m))
     @test_throws ErrorException dot(x[1] * x[2] * x[3], m)
-    @test dot(0.5 * x[1] * x[2]^2, m) == 2.0
-    @test dot(m, x[1] * x[3]) == 3
+    @test (@inferred dot(0.5 * x[1] * x[2]^2, m)) == 2.0
+    @test (@inferred dot(m, x[1] * x[3])) == 3
 end