Clean up matchingvalue etc

Author: penelopeysm

HISTORY.md

@@ -1,5 +1,13 @@
 # DynamicPPL Changelog
 
+## 0.39.10
+
+Rename the internal functions `matchingvalue` and `get_matching_type` to `convert_model_argument` and `promote_model_type_argument` respectively.
+The behaviour of `promote_model_type_argument` has also been slightly changed in some edge cases: for example, `promote_model_type_argument(ForwardDiff.Dual{Nothing,Float64,0}, Vector{Real})` now returns `Vector{ForwardDiff.Dual{Nothing,Real,0}}` instead of `Vector{ForwardDiff.Dual{Nothing,Float64,0}}`.
+In other words, abstract types in the type argument are now properly respected.
+
+This should have almost no impact on end users (unless you were passing `::Type{T}=Vector{Real}` into the model, with an abstract eltype).
+
 ## 0.39.9
 
 The internals of `LogDensityFunction` have been changed slightly so that you do not need to specify `function_annotation` when performing AD with Enzyme.jl.

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.39.9"
+version = "0.39.10"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/compiler.jl

@@ -629,7 +629,6 @@ function add_return_to_last_statment(body::Expr)
     return Expr(body.head, new_args...)
 end
 
-const FloatOrArrayType = Type{<:Union{AbstractFloat,AbstractArray}}
 hasmissing(::Type) = false
 hasmissing(::Type{>:Missing}) = true
 hasmissing(::Type{<:AbstractArray{TA}}) where {TA} = hasmissing(TA)
@@ -754,54 +753,76 @@ function warn_empty(body)
     return nothing
 end
 
-# TODO(mhauru) matchingvalue has methods that can accept both types and values. Why?
-# TODO(mhauru) This function needs a more comprehensive docstring.
 """
-    matchingvalue(param_eltype, value)
-
-Convert the `value` to the correct type, given the element type of the parameters
-being used to evaluate the model.
-"""
-function matchingvalue(param_eltype, value)
-    T = typeof(value)
-    if hasmissing(T)
-        _value = convert(get_matching_type(param_eltype, T), value)
-        # TODO(mhauru) Why do we make a deepcopy, even though in the !hasmissing branch we
-        # are happy to return `value` as-is?
-        if _value === value
-            return deepcopy(_value)
+    convert_model_argument(param_eltype, model_argument)
+
+Convert `model_argument` to the correct type, given the element type of the parameters being
+used to evaluate the model. This function potentially also deep-copies `model_argument` if it
+contains `missing` values.
+"""
+function convert_model_argument(param_eltype, model_argument)
+    T = typeof(model_argument)
+    # If the argument contains missing data, then we potentially need to deepcopy it. This
+    # is because the argument may be e.g. a vector of missings, and evaluating a
+    # tilde-statement like x[1] ~ Normal() would set x[1] = some_not_missing_value, thus
+    # mutating x. If you then run the model again with the same argument, x[1] would no
+    # longer be missing.
+    return if hasmissing(T)
+        # It is possible that we could skip the deepcopy, if the argument has to be promoted
+        # anyway. For example, if we are running with ForwardDiff and the argument is a
+        # Vector{Union{Missing, Float64}}, then we will convert it to a
+        # Vector{Union{Missing, ForwardDiff.Dual{...}}} anyway, which will avoid mutating
+        # the original argument. We can check for this by first converting and then only
+        # deepcopying if the converted value aliases the original.
+        converted_argument = convert(
+            promote_model_type_argument(param_eltype, T), model_argument
+        )
+        if converted_argument === model_argument
+            deepcopy(model_argument)
         else
-            return _value
+            converted_argument
         end
     else
-        return value
+        model_argument
     end
 end
-
-function matchingvalue(param_eltype, value::FloatOrArrayType)
-    return get_matching_type(param_eltype, value)
-end
-function matchingvalue(param_eltype, ::TypeWrap{T}) where {T}
-    return TypeWrap{get_matching_type(param_eltype, T)}()
-end
-
-# TODO(mhauru) This function needs a more comprehensive docstring. What is it for?
-"""
-    get_matching_type(param_eltype, ::TypeWrap{T}) where {T}
-
-Get the specialized version of type `T`, given an element type of the parameters
-being used to evaluate the model.
-"""
-get_matching_type(_, ::Type{T}) where {T} = T
-function get_matching_type(param_eltype, ::Type{<:Union{Missing,AbstractFloat}})
-    return Union{Missing,float_type_with_fallback(param_eltype)}
-end
-function get_matching_type(param_eltype, ::Type{<:AbstractFloat})
-    return float_type_with_fallback(param_eltype)
-end
-function get_matching_type(param_eltype, ::Type{<:Array{T,N}}) where {T,N}
-    return Array{get_matching_type(param_eltype, T),N}
-end
-function get_matching_type(param_eltype, ::Type{<:Array{T}}) where {T}
-    return Array{get_matching_type(param_eltype, T)}
+# These methods handle arguments that are types rather than values.
+function convert_model_argument(param_eltype, t::Type{<:Union{Real,AbstractArray}})
+    return promote_model_type_argument(param_eltype, t)
+end
+function convert_model_argument(param_eltype, ::TypeWrap{T}) where {T}
+    return TypeWrap{promote_model_type_argument(param_eltype, T)}()
+end
+# If the parameter element type is `Any`, then we don't need to do any conversion (but we
+# might need to deepcopy).
+function convert_model_argument(::Type{Any}, model_argument::T) where {T}
+    return hasmissing(T) ? deepcopy(model_argument) : model_argument
+end
+# Extra methods to avoid method ambiguity.
+convert_model_argument(::Type{Any}, t::Type{<:Union{Real,AbstractArray}}) = t
+convert_model_argument(::Type{Any}, t::TypeWrap{T}) where {T} = t
+
+"""
+    promote_model_type_argument(param_eltype, ::Type{T}) where {T}
+    promote_model_type_argument(param_eltype, ::TypeWrap{T}) where {T}
+
+For arguments to a model that are types rather than values, promote the type `T` to
+match the element type of the parameters being used to evaluate the model.
+"""
+promote_model_type_argument(_, ::Type{T}) where {T} = T
+function promote_model_type_argument(param_eltype, ::Type{T}) where {T<:Real}
+    # TODO(penelopeysm): This actually might still be over-aggressive. For example, if
+    # `param_eltype` is `Float32` and `T` is `Vector{Int}`, then (after going through the
+    # Array method) we will promote to `Vector{Float64}`, which seems unnecessary. However,
+    # there's no way to actually check if `T` is the type of something that will later be
+    # assigned to, so this is 'safe'.
+    return Base.promote_type(param_eltype, T)
+end
+# NOTE(penelopeysm): This doesn't work with other types of AbstractArray. To get around
+# that, one could in principle use ArrayInterface.promote_eltype. However, it doesn't seem
+# like there is (1) demand for that; and (2) sufficiently strong adoption of ArrayInterface
+# to make that worth adding as a dependency.
+function promote_model_type_argument(param_eltype, ::Type{Array{T,N}}) where {T,N}
+    promoted_eltype = promote_model_type_argument(param_eltype, eltype(T))
+    return Array{promoted_eltype,N}
 end

src/model.jl

@@ -1010,12 +1010,14 @@ Return the arguments and keyword arguments to be passed to the evaluator of the
     unwrap_args = [
         if is_splat_symbol(var)
             :(
-                $matchingvalue(
+                $convert_model_argument(
                     $get_param_eltype(varinfo, model.context), model.args.$var
                 )...
             )
         else
-            :($matchingvalue($get_param_eltype(varinfo, model.context), model.args.$var))
+            :($convert_model_argument(
+                $get_param_eltype(varinfo, model.context), model.args.$var
+            ))
         end for var in argnames
     ]
     return quote

test/compiler.jl

@@ -836,4 +836,47 @@ module Issue537 end
         @test vi isa VarInfo
         @test vi[@varname(m)] isa Real
     end
+
+    @testset "convert_model_argument" begin
+        tdual = ForwardDiff.Dual{Nothing,Float64,0}
+        # no-op
+        @test DynamicPPL.convert_model_argument(Float64, 1.0) == 1.0
+        @testset "shouldn't promote types of value arguments" begin
+            # i.e. this shouldn't become a dual.
+            @test DynamicPPL.convert_model_argument(tdual, 1.0) == 1.0
+        end
+        @testset "arrays" begin
+            # convert_model_argument should make sure to not deepcopy arrays if not needed
+            x = [1.0]
+            @test DynamicPPL.convert_model_argument(Float64, x) === x
+            # but if there's a missing in the array, it should
+            y = [1.0, missing]
+            y_converted = DynamicPPL.convert_model_argument(Float64, y)
+            @test y_converted !== y
+            @test isequal(y_converted, y)
+        end
+        @testset "type arguments" begin
+            # These tests with types / TypeWrap as the second argument also test
+            # `promote_model_type_argument`.
+            function test_type_conversion(
+                ::Type{input}, ::Type{target}
+            ) where {input,target}
+                converted_type = DynamicPPL.convert_model_argument(tdual, input)
+                @test converted_type == target
+                typewrap = DynamicPPL.TypeWrap{input}()
+                converted_typewrap = DynamicPPL.convert_model_argument(tdual, typewrap)
+                @test converted_typewrap == DynamicPPL.TypeWrap{target}()
+            end
+            test_type_conversion(Float64, tdual)
+            test_type_conversion(Real, ForwardDiff.Dual{Nothing,Real,0})
+            test_type_conversion(Vector{Float64}, Vector{tdual})
+            test_type_conversion(Vector{Real}, Vector{ForwardDiff.Dual{Nothing,Real,0}})
+            test_type_conversion(Matrix{Float64}, Matrix{tdual})
+            test_type_conversion(Matrix{Real}, Matrix{ForwardDiff.Dual{Nothing,Real,0}})
+            test_type_conversion(Vector{Vector{Float64}}, Vector{Vector{tdual}})
+            test_type_conversion(
+                Vector{Vector{Real}}, Vector{Vector{ForwardDiff.Dual{Nothing,Real,0}}}
+            )
+        end
+    end
 end