Merge pull request #113 from tpapp/tp/fix-staticarray-typeinference

Fix static array type inference.

Author: tpapp

Project.toml

@@ -1,7 +1,7 @@
 name = "TransformVariables"
 uuid = "84d833dd-6860-57f9-a1a7-6da5db126cff"
 authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
-version = "0.8.5"
+version = "0.8.6"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

src/TransformVariables.jl

@@ -7,7 +7,7 @@ using LogExpFunctions
 using LinearAlgebra: UpperTriangular, logabsdet
 using UnPack: @unpack
 using Random: AbstractRNG, GLOBAL_RNG
-using StaticArrays: MMatrix, SMatrix, SArray
+using StaticArrays: MMatrix, SMatrix, SArray, SVector, pushfirst
 
 import ChangesOfVariables
 import InverseFunctions

src/aggregation.jl

@@ -145,16 +145,22 @@ end
 result_size(::StaticArrayTransformation{D,S}) where {D,S} = fieldtypes(S)
 
 function transform_with(flag::LogJacFlag, transformation::StaticArrayTransformation{D,S},
-                        x::AbstractVector{T}, index) where {D,S,T}
+                        x::AbstractVector{T}, index::Int) where {D,S,T}
     @unpack inner_transformation = transformation
-    ℓ = logjac_zero(flag, robust_eltype(x))
-    SArray{S}(begin
-                  y, ℓΔ, index′ = transform_with(flag, inner_transformation, x, index)
-                  index = index′
-                  ℓ += ℓΔ
-                  y
-              end
-              for _ in 1:D), ℓ, index
+    # NOTE this is a fix for #112, enforcing types taken from the transformation of the
+    # first element.
+    y1, ℓ1, index1 = transform_with(flag, inner_transformation, x, index)
+    L = typeof(ℓ1)
+    let ℓ::L = ℓ1, index::Int = index1
+        function _f(_)
+            y, ℓΔ, index′ = transform_with(flag, inner_transformation, x, index)
+            index = index′
+            ℓ = ℓ + ℓΔ
+            y
+        end
+        yrest = SVector{D-1}(_f(i) for i in 2:D)
+        SArray{S}(pushfirst(yrest, y1)), ℓ, index
+    end
 end
 
 function inverse_eltype(transformation::Union{ArrayTransformation,StaticArrayTransformation},

src/generic.jl

@@ -44,7 +44,7 @@ $(SIGNATURES)
 
 Initial value for log Jacobian calculations.
 """
-logjac_zero(::LogJac, T::Type{<:Real}) = log(one(T))
+logjac_zero(::LogJac, ::Type{T}) where {T<:Real} = log(one(T))
 
 logjac_zero(::NoLogJac, _) = NOLOGJAC
 

test/runtests.jl

@@ -6,7 +6,7 @@ using LogDensityProblems: logdensity, logdensity_and_gradient
 using LogDensityProblemsAD
 using TransformVariables:
     AbstractTransform, ScalarTransform, VectorTransform, ArrayTransformation,
-    unit_triangular_dimension, logistic, logistic_logjac, logit, inverse_and_logjac
+    unit_triangular_dimension, logistic, logistic_logjac, logit, inverse_and_logjac, NOLOGJAC, transform_with
 import ChangesOfVariables, InverseFunctions
 using Enzyme: autodiff, Reverse, Active, Const
 
@@ -640,3 +640,7 @@ end
             d = as(SVector{2}, asℝ₊)))
     @test [domain_label(t, i) for i in 1:dimension(t)] == [".a", ".b[1,1]", ".c[1]", ".d[1]", ".d[2]"]
 end
+
+@testset "static arrays inference" begin
+    @test @inferred transform_with(NOLOGJAC, as(SVector{3, Float64}), zeros(3), 1) == (SVector(0.0, 0.0, 0.0), NOLOGJAC, 4)
+end