VNT Part 2: Solving the issue of block elements #1180
Conversation
Benchmark Report
Computer InformationBenchmark Results |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## mhauru/vnt-for-fastldf #1180 +/- ##
=========================================================
Coverage ? 80.13%
=========================================================
Files ? 42
Lines ? 4298
Branches ? 0
=========================================================
Hits ? 3444
Misses ? 854
Partials ? 0 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
DynamicPPL.jl documentation for PR #1180 is available at: |
mhauru
left a comment
mhauru
left a comment
There was a problem hiding this comment.
Tests pass, and I think this is ready for review. The only thing I think is missing is the support for (concretised) Colons, but I'll put that in a separate PR.
| # whether it's linked | ||
| is_linked::Bool | ||
| # original size of the variable before vectorisation | ||
| original_size::T |
There was a problem hiding this comment.
It's a bit unfortunate to have to create this field, because now you may end up with abstractly typed RangeAndLinked, which obscures the fact that the two fields you really care about, namely range and is_linked, are still concrete. The reason this is needed is that VNT needs to know how much "space" in a PartialArray an instance of RangeAndLinked takes. An alternative to this could be something like giving setindex!!(::VarNamedTuple, ...) an extra kwarg of something like ignore_size_checks.
| # vr.varname_ranges[vn] may infer to have type `Any`. In this case it is helpful to | ||
| # assert that it is a RangeAndLinked, because even though it remains non-concrete, | ||
| # it'll allow the compiler to infer the types of `range` and `is_linked`. | ||
| return vr.varname_ranges[vn]::RangeAndLinked |
There was a problem hiding this comment.
Without this assertion, this model:
@model function demo_one_variable_multiple_constraints(
::Type{TV}=Vector{Float64}
) where {TV}
x = TV(undef, 5)
x[1] ~ Normal()
x[2] ~ InverseGamma(2, 3)
x[3] ~ truncated(Normal(), -5, 20)
x[4:5] ~ Dirichlet([1.0, 2.0])
return (x=x,)
endfails the test that checks that the return type of LogDensityProblems.logdensity(ldf, x) can be inferred. It infers as Any. This happens because the VNT gets a PartialArray where three elements are RangeAndLinked{Tuple{}} and one is a ArrayLikeBlock{RangeAndLinked{Tuple{Int}}, Tuple{Int}}. That makes the element type Any, and _get_range_and_linked's return type infers as Any, and all is lost.
Now, I'm fine saying that a model that mixes x[i] and x[j:l] for the same x is pretty type unstable (note that the priors have to mix univariate and multivariate). However, I don't think that should be licence for even the logdensity return type to infer as Any, and it didn't before this PR. Hence this extra assertion, to have the type instability not be so egregious.
Note that this problem is independent of the above comment of introducing a type parameter to RangeAndLinked: Even without the type parameter, the wrapping in ArrayLikeBlock would force the PartialArray to have element type Any.
| _has_colon(::T) where {T<:Tuple} = any(x <: Colon for x in T.parameters) | ||
| function _is_multiindex(::T) where {T<:Tuple} | ||
| return any(x <: UnitRange || x <: Colon for x in T.parameters) | ||
| # The non-generated function implementations of these would be |
There was a problem hiding this comment.
Without this, getindex(vnt, @varname(y.z[2:3, 3, 2:3])) was type stable but getindex(vnt, @varname(y.z[2, 2:3, 3, 2:3, 4])) was not.
This PR adds a wrapper type
ArrayLikeBlock, thatPartialArrayuses to store any elements that have non-zero size but are notAbstractArrays, when such a value is being set.Other improvements to VNT are also included, that I happened to need and implement while working on this.
This PR is just completing the work in #1150, but I keep it separate to make reviewing easier.