Error informatively on partial specification of a multivariate variable (#2239)

HISTORY.md

@@ -1,3 +1,9 @@
+# 0.42.2
+
+`condition`, `fix`, and `predict` now raise an informative error when given only part of a variable that the model samples as a single multivariate draw (e.g. `x[1:10]` when `x ~ MvNormal(zeros(20), ...)`).
+
+This is not a behaviour change: the partial case never worked, previously failing silently or with an opaque `DimensionMismatch`. To handle components individually, use a loop: `for i in eachindex(x); x[i] ~ ...; end`. See [Turing.jl#2239](https://github.com/TuringLang/Turing.jl/issues/2239).
+
 # 0.42.1
 
 Fixed a type-inference failure that made nested submodels (a `~ to_submodel(...)` statement inside a model that is itself evaluated as a submodel) very slow.

Project.toml

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

ext/DynamicPPLMCMCChainsExt.jl

@@ -273,6 +273,15 @@ If `include_all` is `false`, the returned `Chains` will contain only those varia
 the samples in `chain`. This is useful when you want to sample only new variables from the posterior
 predictive distribution.
 
+!!! warning "Variables are treated as they occur in the model"
+    A variable drawn from a multivariate distribution in a single tilde-statement
+    (e.g. `x ~ MvNormal(...)` or `x ~ filldist(Normal(), n)`) is a *single* random
+    variable, not a collection of i.i.d. components. `predict` cannot fill in a subset of
+    such a variable's components: if `chain` supplies only some of them, an error is
+    raised. To treat components individually, declare them in a loop, e.g.
+    `for i in eachindex(x); x[i] ~ Normal(); end`.
+    See [TuringLang/Turing.jl#2239](https://github.com/TuringLang/Turing.jl/issues/2239).
+
 # Examples
 ```jldoctest
 using AbstractMCMC, Distributions, DynamicPPL, Random

src/compiler.jl

@@ -127,6 +127,53 @@ function contextual_isfixed(context::AbstractContext, vn)
     end
 end
 
+"""
+    _check_supplied_shape(dist, supplied, vn)
+
+Return `supplied`, but error (see TuringLang/Turing.jl#2239) if it specifies only part of a
+variable that the model samples as a single multivariate draw (e.g. `x[1:10]` for
+`x ~ MvNormal(zeros(20), ...)`). `supplied` is either a materialised value (`condition`
+/`fix`) or the parameter `VarNamedTuple` (`predict`/`InitFromParams`); non-multivariate
+`dist`s are a no-op.
+"""
+_check_supplied_shape(_, supplied, _) = supplied
+function _check_supplied_shape(
+    dist::Distributions.MultivariateDistribution, supplied::AbstractArray, vn
+)
+    length(supplied) == length(dist) ||
+        _throw_incomplete_multivariate(vn, length(supplied), dist)
+    return supplied
+end
+function _check_supplied_shape(
+    dist::Distributions.MultivariateDistribution, params::VarNamedTuples.VarNamedTuple, vn
+)
+    # Only a reference to the *whole* variable can be partial; an absent indexed leaf
+    # (`x[i]`) is simply a new variable to be sampled from the prior.
+    AbstractPPL.getoptic(vn) isa AbstractPPL.Iden || return params
+    sym = AbstractPPL.getsym(vn)
+    haskey(params.data, sym) || return params
+    val = params.data[sym]
+    if val isa VarNamedTuples.PartialArray
+        n = count(val.mask)
+        # `n == 0` means nothing was supplied (the caller falls back to the prior); a full,
+        # correctly-sized array would already have been found by `hasvalue`.
+        (n == 0 || n == length(dist)) || _throw_incomplete_multivariate(vn, n, dist)
+    end
+    return params
+end
+@noinline function _throw_incomplete_multivariate(vn, n, dist)
+    return throw(
+        ArgumentError(
+            "A value with $(n) element(s) was supplied for `$(vn)`, but the model samples " *
+            "`$(vn)` as a single $(length(dist))-dimensional random variable. Supplying a " *
+            "subset of such a variable (via `condition`, `fix`, or `predict` with a chain " *
+            "from a differently-sized model) is not supported. To treat its components " *
+            "individually, declare them in a loop, e.g. " *
+            "`for i in eachindex($(vn)); $(vn)[i] ~ ...; end`.",
+        ),
+    )
+end
+
 # If we're working with, say, a `Symbol`, then we're not going to `view`.
 maybe_view(x) = x
 maybe_view(x::Expr) = :(@views($x))
@@ -486,8 +533,12 @@ function generate_tilde(left, right)
             # need to use Accessors.set to safely set it.
             $(assign_or_set!!(
                 left,
-                :($(DynamicPPL.getfixed_nested)(
-                    __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
+                :($(DynamicPPL._check_supplied_shape)(
+                    $dist,
+                    $(DynamicPPL.getfixed_nested)(
+                        __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
+                    ),
+                    $vn,
                 )),
                 vn,
             ))
@@ -505,8 +556,12 @@ function generate_tilde(left, right)
             $supplied_val = if $(DynamicPPL.inargnames)($vn, __model__)
                 $(maybe_view(left))
             else
-                $(DynamicPPL.getconditioned_nested)(
-                    __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
+                $(DynamicPPL._check_supplied_shape)(
+                    $dist,
+                    $(DynamicPPL.getconditioned_nested)(
+                        __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
+                    ),
+                    $vn,
                 )
             end
 

src/contexts/init.jl

@@ -188,6 +188,8 @@ function init(
             TransformedValue(x, NoTransform())
         end
     else
+        # Error if only part of a whole multivariate variable was supplied (#2239).
+        _check_supplied_shape(dist, p.params, vn)
         p.fallback === nothing && error("No value was provided for the variable `$(vn)`.")
         init(rng, vn, dist, p.fallback)
     end

src/model.jl

@@ -522,6 +522,14 @@ Return a `Model` which now treats the variables in `values` as fixed.
 
 See also: [`unfix`](@ref), [`fixed`](@ref)
 
+!!! warning "Fixing applies to whole variables"
+    Variables are treated as they occur in the model. A variable drawn from a multivariate
+    distribution in a single tilde-statement (e.g. `x ~ MvNormal(...)`) is a *single* random
+    variable, so a subset of its components cannot be fixed independently; attempting to do
+    so raises an error. Declare components in a loop (`x[i] ~ ...`) if you need to fix them
+    individually.
+    See [TuringLang/Turing.jl#2239](https://github.com/TuringLang/Turing.jl/issues/2239).
+
 # Examples
 ## Simple univariate model
 ```jldoctest fix

test/conditionfix.jl

@@ -541,6 +541,20 @@ end
     end
 end
 
+@testset "partial multivariate variable errors (#2239)" begin
+    @model function mvc(n)
+        m ~ MvNormal(zeros(n), 1.0)
+        return m
+    end
+    # Supplying only part of a variable the model samples as one multivariate draw is
+    # unsupported: `fix` used to silently collapse it and `condition` to throw an opaque
+    # `DimensionMismatch`. Both must now raise the informative #2239 error.
+    @test_throws ArgumentError fix(mvc(3), Dict(@varname(m[1]) => 1.0))()
+    @test_throws ArgumentError condition(mvc(3), Dict(@varname(m[1]) => 1.0))()
+    # A whole-variable value of the correct size is still honoured.
+    @test fix(mvc(3), Dict(@varname(m) => [1.0, 2.0, 3.0]))() == [1.0, 2.0, 3.0]
+end
+
 @info "Completed $(@__FILE__) in $(now() - __now__)."
 
 end # module

test/model.jl

@@ -642,6 +642,30 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
                 pdns = DynamicPPL.predict(fmodel, chn[[:b]])
                 @test Set(keys(pdns)) == Set([:x])
             end
+
+            @testset "errors on partial multivariate variable (#2239)" begin
+                @model function mv_partial(n)
+                    return m ~ MvNormal(zeros(n), 1.0)
+                end
+                # The n=10 chain carries only `m[1:10]`; predicting with the n=20 model
+                # asks to fill part of the single multivariate `m`, which must error
+                # rather than silently resample it.
+                chn10 = make_chain_from_prior(mv_partial(10), 5)
+                @test_throws ArgumentError DynamicPPL.predict(mv_partial(20), chn10)
+            end
+
+            @testset "per-index variables are not falsely flagged (#2239)" begin
+                @model function iid(k)
+                    x = Vector{Float64}(undef, k)
+                    x .~ Normal()
+                    return x
+                end
+                # Each `x[i]` is its own variable, so predicting with a larger model must
+                # not error; the extra indices are simply sampled from the prior.
+                chn10 = make_chain_from_prior(iid(10), 5)
+                pred = DynamicPPL.predict(iid(20), chn10)
+                @test Symbol("x[20]") in keys(pred)
+            end
         end
     end