Allow empty subsets of VarInfos (#692)

* Allow empty subsets of VarInfos * Run JuliaFormatter Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
TuringLang · Oct 17, 2024 · 54691bf · 54691bf · mhauru · Oct 18, 2024
1 parent 27ba772
commit 54691bf
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 17 deletions.
diff --git a/src/simple_varinfo.jl b/src/simple_varinfo.jl
@@ -439,22 +439,17 @@ function subset(varinfo::SimpleVarInfo, vns::AbstractVector{<:VarName})
  return Accessors.@set varinfo.values = _subset(varinfo.values, vns)
 end
 
-function _subset(x::AbstractDict, vns)
+function _subset(x::AbstractDict, vns::AbstractVector{VN}) where {VN<:VarName}
  vns_present = collect(keys(x))
- vns_found = mapreduce(vcat, vns) do vn
+ vns_found = mapreduce(vcat, vns; init=VN[]) do vn
  return filter(Base.Fix1(subsumes, vn), vns_present)
  end
-
- # NOTE: This `vns` to be subsume varnames explicitly present in `x`.
+ C = ConstructionBase.constructorof(typeof(x))
  if isempty(vns_found)
- throw(
- ArgumentError(
- "Cannot subset `AbstractDict` with `VarName` which does not subsume any keys.",
- ),
- )
+ return C()
+ else
+ return C(vn => x[vn] for vn in vns_found)
  end
- C = ConstructionBase.constructorof(typeof(x))
- return C(vn => x[vn] for vn in vns_found)
 end
 
 function _subset(x::NamedTuple, vns)

diff --git a/src/varinfo.jl b/src/varinfo.jl
@@ -368,20 +368,24 @@ function subset(varinfo::TypedVarInfo, vns::AbstractVector{<:VarName})
  )
 end
 
-function subset(metadata::Metadata, vns_given::AbstractVector{<:VarName})
+function subset(metadata::Metadata, vns_given::AbstractVector{VN}) where {VN<:VarName}
  # TODO: Should we error if `vns` contains a variable that is not in `metadata`?
  # For each `vn` in `vns`, get the variables subsumed by `vn`.
- vns = mapreduce(vcat, vns_given) do vn
+ vns = mapreduce(vcat, vns_given; init=VN[]) do vn
  filter(Base.Fix1(subsumes, vn), metadata.vns)
  end
  indices_for_vns = map(Base.Fix1(getindex, metadata.idcs), vns)
- indices = Dict(vn => i for (i, vn) in enumerate(vns))
+ indices = if isempty(vns)
+ Dict{VarName,Int}()
+ else
+ Dict(vn => i for (i, vn) in enumerate(vns))
+ end
  # Construct new `vals` and `ranges`.
  vals_original = metadata.vals
  ranges_original = metadata.ranges
  # Allocate the new `vals`. and `ranges`.
- vals = similar(metadata.vals, sum(length, ranges_original[indices_for_vns]))
- ranges = similar(ranges_original)
+ vals = similar(metadata.vals, sum(length, ranges_original[indices_for_vns]; init=0))
+ ranges = similar(ranges_original, length(vns))
  # The new range `r` for `vns[i]` is offset by `offset` and
  # has the same length as the original range `r_original`.
  # The new `indices` (from above) ensures ordering according to `vns`.
@@ -415,7 +419,7 @@ function subset(metadata::Metadata, vns_given::AbstractVector{<:VarName})
  ranges,
  vals,
  metadata.dists[indices_for_vns],
- metadata.gids,
+ metadata.gids[indices_for_vns],
  metadata.orders[indices_for_vns],
  flags,
  )

diff --git a/test/varinfo.jl b/test/varinfo.jl
@@ -578,6 +578,13 @@ DynamicPPL.getspace(::DynamicPPL.Sampler{MySAlg}) = (:s,)
  else
  vns_supported_standard
  end
+
+ @testset ("$(convert(Vector{VarName}, vns_subset)) empty") for vns_subset in
+ vns_supported
+ varinfo_subset = subset(varinfo, VarName[])
+ @test isempty(varinfo_subset)
+ end
+
  @testset "$(convert(Vector{VarName}, vns_subset))" for vns_subset in
  vns_supported
  varinfo_subset = subset(varinfo, vns_subset)