ML/FIML: workaround generic_matmul issue

Author: alyst

src/loss/ML/FIML.jl

@@ -164,16 +164,18 @@ end
 function ∇F_fiml_outer!(G, JΣ, Jμ, fiml::SemFIML)
     implied = fiml.implied
 
+    I_A⁻¹ = parent(implied.I_A⁻¹)
+    F⨉I_A⁻¹ = parent(implied.F * I_A⁻¹)
+    S = parent(implied.S)
+
     Iₙ = sparse(1.0I, size(implied.A)...)
-    P = kron(implied.F⨉I_A⁻¹, implied.F⨉I_A⁻¹)
-    Q = kron(implied.S * implied.I_A⁻¹', Iₙ)
+    P = kron(F⨉I_A⁻¹, F⨉I_A⁻¹)
+    Q = kron(S * I_A⁻¹', Iₙ)
     Q .+= fiml.commutator * Q
 
     ∇Σ = P * (implied.∇S + Q * implied.∇A)
 
-    ∇μ =
-        implied.F⨉I_A⁻¹ * implied.∇M +
-        kron((implied.I_A⁻¹ * implied.M)', implied.F⨉I_A⁻¹) * implied.∇A
+    ∇μ = F⨉I_A⁻¹ * implied.∇M + kron((I_A⁻¹ * implied.M)', F⨉I_A⁻¹) * implied.∇A
 
     mul!(G, ∇Σ', JΣ) # actually transposed
     mul!(G, ∇μ', Jμ, -1, 1)

src/loss/ML/ML.jl

@@ -184,9 +184,9 @@ function evaluate!(objective, gradient, hessian, ml::SemML, par)
     end
 
     if !isnothing(gradient)
-        S = implied.S
-        F⨉I_A⁻¹ = implied.F⨉I_A⁻¹
-        I_A⁻¹ = implied.I_A⁻¹
+        S = parent(implied.S)
+        F⨉I_A⁻¹ = parent(implied.F⨉I_A⁻¹)
+        I_A⁻¹ = parent(implied.I_A⁻¹)
         ∇A = implied.∇A
         ∇S = implied.∇S
 
@@ -198,15 +198,9 @@ function evaluate!(objective, gradient, hessian, ml::SemML, par)
         C = mul!(
             ml.varXvar_1,
             F⨉I_A⁻¹',
-            mul!(ml.obsXvar_1, Symmetric(mul!(ml.obsXobs_3, one_Σ⁻¹Σₒ, Σ⁻¹)), F⨉I_A⁻¹),
-        )
-        mul!(
-            gradient,
-            ∇A',
-            vec(mul!(ml.varXvar_3, Symmetric(C), mul!(ml.varXvar_2, S, I_A⁻¹'))),
-            2,
-            0,
+            mul!(ml.obsXvar_1, mul!(ml.obsXobs_3, one_Σ⁻¹Σₒ, Σ⁻¹), F⨉I_A⁻¹),
         )
+        mul!(gradient, ∇A', vec(mul!(ml.varXvar_3, C, mul!(ml.varXvar_2, S, I_A⁻¹'))), 2, 0)
         mul!(gradient, ∇S', vec(C), 1, 1)
 
         if MeanStruct(implied) === HasMeanStruct