[Merged by Bors] - Sibling PR of introduction of Setfield.jl in AbstractPPL.jl

Project.toml

@@ -10,6 +10,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]

src/DynamicPPL.jl

@@ -10,6 +10,8 @@ using ChainRulesCore: ChainRulesCore
 using MacroTools: MacroTools
 using ZygoteRules: ZygoteRules
 
+using Setfield: Setfield
+
 using Random: Random
 
 import Base:

src/compiler.jl

@@ -19,7 +19,7 @@ function isassumption(expr::Union{Symbol,Expr})
  vn = gensym(:vn)
 
  return quote
- let $vn = $(varname(expr))
+ let $vn = $(varname(expr, true))
  # This branch should compile nicely in all cases except for partial missing data
  # For example, when `expr` is `:(x[i])` and `x isa Vector{Union{Missing, Float64}}`
  if !$(DynamicPPL.inargnames)($vn, __model__) ||
@@ -38,7 +38,7 @@ isassumption(expr) = :(false)
 
 # If we're working with, say, a `Symbol`, then we're not going to `view`.
 maybe_view(x) = x
-maybe_view(x::Expr) = :($(DynamicPPL.maybe_unwrap_view)(@view($x)))
+maybe_view(x::Expr) = :($(DynamicPPL.maybe_unwrap_view)(@views($x)))
 
 # If the result of a `view` is a zero-dim array then it's just a
 # single element. Likely the rest is expecting type `eltype(x)`, hence
@@ -90,6 +90,28 @@ left-hand side of a `.~` expression such as `x .~ Normal()`.
 
 This is used mainly to unwrap `NamedDist` distributions and adjust the indices of the
 variables.
+
+# Examples
+```jldoctest
+julia> _, _, vns = DynamicPPL.unwrap_right_left_vns(MvNormal(1, 1.0), randn(1, 2), @varname(x)); vns
+2-element Vector{VarName{:x, Setfield.IndexLens{Tuple{Colon, Int64}}}}:
+ x[:,1]
+ x[:,2]
+
+julia> _, _, vns = DynamicPPL.unwrap_right_left_vns(Normal(), randn(1, 2), @varname(x)); vns
+1×2 Matrix{VarName{:x, Setfield.IndexLens{Tuple{Int64, Int64}}}}:
+ x[1,1] x[1,2]
+
+julia> _, _, vns = DynamicPPL.unwrap_right_left_vns(Normal(), randn(1, 2), @varname(x[:])); vns
+1×2 Matrix{VarName{:x, Setfield.ComposedLens{Setfield.IndexLens{Tuple{Colon}}, Setfield.IndexLens{Tuple{Int64, Int64}}}}}:
+ x[:][1,1] x[:][1,2]
+
+julia> _, _, vns = DynamicPPL.unwrap_right_left_vns(Normal(), randn(3), @varname(x[1])); vns
+3-element Vector{VarName{:x, Setfield.ComposedLens{Setfield.IndexLens{Tuple{Int64}}, Setfield.IndexLens{Tuple{Int64}}}}}:
+ x[1][1]
+ x[1][2]
+ x[1][3]
+```
 """
 unwrap_right_left_vns(right, left, vns) = right, left, vns
 function unwrap_right_left_vns(right::NamedDist, left, vns)
@@ -103,7 +125,7 @@ function unwrap_right_left_vns(
  # for `i = size(left, 2)`. Hence the symbol should be `x[:, i]`,
  # and we therefore add the `Colon()` below.
  vns = map(axes(left, 2)) do i
- return VarName(vn, (vn.indexing..., Colon(), Tuple(i)))
+ return vn ∘ Setfield.IndexLens((Colon(), i))
  end
  return unwrap_right_left_vns(right, left, vns)
 end
@@ -113,7 +135,7 @@ function unwrap_right_left_vns(
  vn::VarName,
 )
  vns = map(CartesianIndices(left)) do i
- return VarName(vn, (vn.indexing..., Tuple(i)))
+ return vn ∘ Setfield.IndexLens(Tuple(i))
  end
  return unwrap_right_left_vns(right, left, vns)
 end
@@ -271,6 +293,10 @@ function generate_mainbody!(mod, found, expr::Expr, warn)
  # Do not touch interpolated expressions
  expr.head === :$ && return expr.args[1]
 
+ # Do we don't want escaped expressions because we unfortunately
+ # escape the entire body afterwards.
+ Meta.isexpr(expr, :escape) && return generate_mainbody(mod, found, expr.args[1], warn)
+
  # If it's a macro, we expand it
  if Meta.isexpr(expr, :macrocall)
  return generate_mainbody!(mod, found, macroexpand(mod, expr; recursive=true), warn)
@@ -303,95 +329,161 @@ function generate_mainbody!(mod, found, expr::Expr, warn)
  return Expr(expr.head, map(x -> generate_mainbody!(mod, found, x, warn), expr.args)...)
 end
 
+function generate_tilde_literal(left, right)
+ # If the LHS is a literal, it is always an observation
+ return quote
+ $(DynamicPPL.tilde_observe!)(
+ __context__, $(DynamicPPL.check_tilde_rhs)($right), $left, __varinfo__
+ )
+ end
+end
+
 """
  generate_tilde(left, right)
 
 Generate an `observe` expression for data variables and `assume` expression for parameter
 variables.
 """
 function generate_tilde(left, right)
- # If the LHS is a literal, it is always an observation
- if isliteral(left)
- return quote
- $(DynamicPPL.tilde_observe!)(
- __context__, $(DynamicPPL.check_tilde_rhs)($right), $left, __varinfo__
- )
- end
- end
+ isliteral(left) && return generate_tilde_literal(left, right)
 
  # Otherwise it is determined by the model or its value,
  # if the LHS represents an observation
- @gensym vn inds isassumption
+ @gensym vn isassumption
+
  return quote
- $vn = $(varname(left))
- $inds = $(vinds(left))
- $isassumption = $(DynamicPPL.isassumption(left))
+ $vn = $(remove_escape(varname(left, true)))
+ $isassumption = $(remove_escape(DynamicPPL.isassumption(left)))
  if $isassumption
- $left = $(DynamicPPL.tilde_assume!)(
- __context__,
- $(DynamicPPL.unwrap_right_vn)(
- $(DynamicPPL.check_tilde_rhs)($right), $vn
- )...,
- $inds,
- __varinfo__,
- )
+ $(generate_tilde_assume(left, right, vn))
  else
  $(DynamicPPL.tilde_observe!)(
  __context__,
  $(DynamicPPL.check_tilde_rhs)($right),
  $(maybe_view(left)),
  $vn,
- $inds,
  __varinfo__,
  )
  end
  end
 end
 
+function generate_tilde_assume(left::Symbol, right, vn)
+ return quote
+ $left = $(DynamicPPL.tilde_assume!)(
+ __context__,
+ $(DynamicPPL.unwrap_right_vn)(
+ $(DynamicPPL.check_tilde_rhs)($right), $vn
+ )...,
+ __varinfo__,
+ )
+ end
+end
+
+function generate_tilde_assume(left::Expr, right, vn)
+ expr = :(
+ $left = $(DynamicPPL.tilde_assume!)(
+ __context__,
+ $(DynamicPPL.unwrap_right_vn)(
+ $(DynamicPPL.check_tilde_rhs)($right), $vn
+ )...,
+ __varinfo__,
+ )
+ )
+
+ return remove_escape(setmacro(identity, expr, overwrite=true))
+end
+
 """
  generate_dot_tilde(left, right)
 
 Generate the expression that replaces `left .~ right` in the model body.
 """
 function generate_dot_tilde(left, right)
- # If the LHS is a literal, it is always an observation
- if isliteral(left)
- return quote
- $(DynamicPPL.dot_tilde_observe!)(
- __context__, $(DynamicPPL.check_tilde_rhs)($right), $left, __varinfo__
- )
- end
- end
+ isliteral(left) && return generate_tilde_literal(left, right)
 
  # Otherwise it is determined by the model or its value,
  # if the LHS represents an observation
- @gensym vn inds isassumption
+ @gensym vn isassumption
  return quote
- $vn = $(varname(left))
- $inds = $(vinds(left))
+ $vn = $(varname(left, true))
  $isassumption = $(DynamicPPL.isassumption(left))
  if $isassumption
- $left .= $(DynamicPPL.dot_tilde_assume!)(
- __context__,
- $(DynamicPPL.unwrap_right_left_vns)(
- $(DynamicPPL.check_tilde_rhs)($right), $(maybe_view(left)), $vn
- )...,
- $inds,
- __varinfo__,
- )
+ $(generate_dot_tilde_assume(left, right, vn))
  else
  $(DynamicPPL.dot_tilde_observe!)(
  __context__,
  $(DynamicPPL.check_tilde_rhs)($right),
  $(maybe_view(left)),
  $vn,
- $inds,
  __varinfo__,
  )
  end
  end
 end
 
+function generate_dot_tilde_assume(left::Symbol, right, vn)
+ return :(
+ $left .= $(DynamicPPL.dot_tilde_assume!)(
+ __context__,
+ $(DynamicPPL.unwrap_right_left_vns)(
+ $(DynamicPPL.check_tilde_rhs)($right), $(maybe_view(left)), $vn
+ )...,
+ __varinfo__,
+ )
+ )
+end
+
+function generate_dot_tilde_assume(left::Expr, right, vn)
+ expr = :(
+ $left .= $(DynamicPPL.dot_tilde_assume!)(
+ __context__,
+ $(DynamicPPL.unwrap_right_left_vns)(
+ $(DynamicPPL.check_tilde_rhs)($right), $(maybe_view(left)), $vn
+ )...,
+ __varinfo__,
+ )
+ )
+
+ return remove_escape(setmacro(identity, expr, overwrite=true))
+end
+
+# HACK: This is unfortunate. It's a consequence of the fact that in
+# DynamicPPL we the entire function body. Instead we should be
+# more selective with our escape. Until that's the case, we remove them all.
+remove_escape(x) = x
+function remove_escape(expr::Expr)
+ Meta.isexpr(expr, :escape) && return remove_escape(expr.args[1])
+ return Expr(expr.head, map(x -> remove_escape(x), expr.args)...)
+end
+
+# TODO: Make PR to Setfield.jl to use `gensym` for the `lens` variable.
+# This seems like it should be the case anyways since it allows multiple
+# calls to `setmacro` without any cost to the current functionality.
+function setmacro(lenstransform, ex::Expr; overwrite::Bool=false)
+ @assert ex.head isa Symbol
+ @assert length(ex.args) == 2
+ ref, val = ex.args
+ obj, lens = Setfield.parse_obj_lens(ref)
+ lens_var = gensym("lens")
+ dst = overwrite ? obj : gensym("_")
+ val = esc(val)
+ ret = if ex.head == :(=)
+ quote
+ $lens_var = ($lenstransform)($lens)
+ $dst = $(Setfield.set)($obj, $lens_var, $val)
+ end
+ else
+ op = get_update_op(ex.head)
+ f = :($(Setfield._UpdateOp)($op,$val))
+ quote
+ $lens_var = ($lenstransform)($lens)
+ $dst = $(Setfield.modify)($f, $obj, $lens_var)
+ end
+ end
+ ret
+end
+
 const FloatOrArrayType = Type{<:Union{AbstractFloat,AbstractArray}}
 hasmissing(T::Type{<:AbstractArray{TA}}) where {TA<:AbstractArray} = hasmissing(TA)
 hasmissing(T::Type{<:AbstractArray{>:Missing}}) = true