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.
There was a problem hiding this comment.
yes, either that -- or you could just pass the size itself as a kwarg and if the kwarg is not passed, it can fall back to calling vnt_size on the value being set. not sure if that is bad for type stability as i've been bitten by julia kwargs many times but from an interface perspective, i think that is most flexible as it allows callers to specifically override the auto-size determination mechanism. my suspicion is type stability should be ok though since it will become Core.kwcall with NamedTuple{size::Tuple{Int,Int}} (or whatever the size is) which is all concretely typed.
when constructing the VNT of RangeAndLinked (inside the LDF setup code) it should be trivial to determine the size separately from the range + linked status and pass it into setindex. right now we would just need to read from the Metadata dist field, otherwise if it's an accumulator then the info is passed through from accumulate_assume!!. after the VNT has been constructed we would only ever read from the VNT, so there would no longer be any need to know the size and it can be dropped from this struct
There was a problem hiding this comment.
I agree that it makes this extra unfortunate that this is only needed when setting the value, by getindex-time it's not needed. I'm not necessarily opposed to something like the kwarg solution, though I'm not a big fan of how it's different from BangBang.setindex!! for everything else. I simultaneously hold these two conflicting opinions:
- If some library defines a function
fthat operates for many types, and it uses the (semantically) same arguments for almost all of them, then if to define your own method forfyou need to change the arguments, that means you shouldn't define the thing your defining as a method off, but rather as a different function. setindex!!is clearly the right function to define a method for for VNT, and anything else would be an unnecessary convolution of the interface that makes VNT less intuitive and more laborious to use.
I don't think it's an unreasonable request to say that if you want to set a value of type X into an array as a block, so that it covers some range like [1:3, 4:17, 100:101], then type X must have some suitable notion of size or otherwise it's a no-go. You could say that the notion of size something you provide ephemerally at setindex!! time, but it feels a bit hacky, compared to saying that X actually "owns" its notion of size.
There was a problem hiding this comment.
That makes sense. I think, the purist part of me (which is quite a large part of me...) much prefers your first argument over your second.
Why not both, though? We could have BangBang.setindex!! for 99% of use cases, and something like DynamicPPL.setindex!!_sized that allows you to pass the size (or skip the check and infer from the varname; I'm ok with either). The implementation would be mostly shared
BangBang.setindex!!(..., val) = DynamicPPL.setindex!!_sized(..., val, vnt_size(val))so code duplication would not be too much; and presumably one can explain in the docstring that the latter should only be used when you the programmer promises that it's a sensible operation (maybe call it unsafe or anything, but it's not really unsafe since there's no real danger in setting a wrong size, as long as you are always consistent with the indices used for getting and setting, which the code enforces).
I guess it is a bit more complexity; personally I would be OK with making that tradeoff in return for the satisfaction of knowing that we have done Good Programming Practices.
There was a problem hiding this comment.
Another way of framing it is that we're going to have to introduce a hack somewhere to make this work, and I think I prefer the hack of controlling the size (or equivalently skipping the check and inferring it from the varname key; come to think of it, I'm starting to prefer that option) vs. the hack of bundling additional fields into a data structure that doesn't really need it. I don't think that RangeAndLinked intrinsically needs the size: if we were using a Dict{VarName} then there's no need for the size (although of course that comes with other drawbacks). So the way I see it, the reason for needing a hack is VNT, and it seems unfair to force RangeAndLinked to adjust itself for someone else.
| # 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.
There was a problem hiding this comment.
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
I think in principle we should be able to do better (but maybe not in practice). Once you remove the Tuple from RangeAndLinked, the definition of _getindex on a PartialArray{Union{T,ArrayLikeBlock{T}}} suggests that it can only ever return a T or error, and should thus in theory be type stable as long as T is concrete. Of course there are two problems with this. One is assuming that Julia's type inference is capable of performing this analysis. The other is assuming that the eltype will be Union{T,ArrayLikeBlock{T}} rather than Any, which I think will depend on the semantics of things like typejoin that are beyond our control.
I'm inclined to say that if Julia is capable of inferring a return type of T from _getindex(::PartialArray{Union{T,ArrayLikeBlock{T}}}), then it may or may not become worthwhile to use our own version of typejoin that special-cases Union{T,ArrayLikeBlock{T}}, a bit like how nothing / missing is special-cased right now. But that can be investigated at another time, happy to keep the assertion right now.
There was a problem hiding this comment.
This, too, I think is an excellent idea. The problem is indeed typejoin, but I haven't thought about special casing ArrayLikeBlock like that.
| _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.
| original_size::T | ||
| end | ||
|
|
||
| Base.size(ral::RangeAndLinked) = ral.original_size |
There was a problem hiding this comment.
Not looking at the whole pr, too late, I've looked -- but wanted to comment on this line -- If one were to look at this line without all the VNT stuff context, I think it could easily look illogical -- firstly there is no broader purpose for defining Base.size on RangeAndLinked (conceptually RangeAndLinked is not a container or AbstractArray, and the range is not meant to be iterated over); secondly even if there was, this is not an unambiguously 'correct' definition of Base.size as opposed to say size(ral.range).
It appears to me that this definition only exists to satisfy the VNT interface, and for this reason I think it is better to not rely on Base.size but instead introduce a separate interface function -- call it vnt_size or anything -- that is specifically used for setting blocks in VNTs.
Using a separate function also avoids issues with structs imported from other libraries where Base.size is either not defined (and defining it would be type piracy); or is defined in a way that is not what we want it to be for VNT (we can't override the definition). As a concrete example, Base.size(::ProductNamedTupleDistribution) does not exist:
julia> size(product_distribution((; a = Normal())))
ERROR: MethodError: no method matching size(::Distributions.ProductNamedTupleDistribution{(:a,), Tuple{Normal{Float64}}, Continuous, Float64})We could of course still have a fallback definition of
vnt_size(x) = size(x)to avoid redefining size on everything.
There was a problem hiding this comment.
I think this is a very good idea. Will implement, not sure if I'll do it now or in January.
| # If _setindex!! works correctly, we should only be able to reach this point if all | ||
| # the elements in `val` are identical to first_elem. Thus we just return that one. | ||
| return first(val).block |
| end | ||
| @generated function _is_multiindex(::T) where {T<:Tuple} | ||
| for x in T.parameters | ||
| if x <: UnitRange || x <: Colon |
There was a problem hiding this comment.
is UnitRange too specific here? i think it should be AbstractVector{<:Integer} since you can in principle write e.g. x[1:1:3] = rand(3) or even x[[2, 3, 456]] = rand(3) -- RangeAndLinked assumes contiguity but that is true by construction.
There was a problem hiding this comment.
In one of the later PRs I change it to be AbstractUnitRange, which I didn't know was a thing. Includes things likeBase.OneTo.
In January when we start reviewing this properly, I wonder if what we should do is take the last PR in this VNT series (whatever that ends up being at that time), take the version of varnamedtuple.jl in it (both in src and in test) and make a separate PR that just merges those files into breaking. You could then review those files in their final state, rather than have to go through the chronology of how I gradually improved them. Then we would merge that into all these other PRs, and they would be left with just a diff that makes the changes in the rest of DPPL to use the new data structure. Anyway, that's for January.
There was a problem hiding this comment.
True, but in principle you could have x[[2, 4, 6]] ~ Dirichlet(ones(3)) and I thinkkkk there's no reason why we need to forbid that as your code should handle it perfectly fine apart from this type bound.
I'm not very fussed about how to review; happy to chat about what would work best for you when new year rolls round (I doubt I will be looking at this tomorrow, since the piano arrived 😄). I like having the separate PRs just because it's easier to fit it into my brain bit by bit, so I might ask that you not get rid of them, but maybe you could have a single big PR to breaking and then make any changes that we agree on to that branch. I did find it quite annoying to continually handle the rebase conflicts with the InitContext series of PRs, especially for the first few ones because I had to propagate the changes up through several later PRs.
There was a problem hiding this comment.
As far as this PR is concerned though, I'm quite happy with merging apart from the stuff above so feel free to go ahead if you are feeling particularly productive today haha
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.