Loop with condition inside fails?

[gdalle]

Am I nesting the @trace macros wrong?

julia> using Reactant

julia> function myfunc_traced(x)  # compute sum of positive elements in x
           s = zero(eltype(x))
           @trace for i in eachindex(x)
               @allowscalar cond = x[i] > 0
               @trace if cond
                   @allowscalar s += x[i]
               end
           end
           return s
       end
ERROR: LoadError: Cannot use @trace on a block that contains a return statement
Stacktrace:
   [1] error(s::String)
     @ Base ./error.jl:35
   [2] #28
     @ ~/.julia/packages/ReactantCore/SEqVX/src/ReactantCore.jl:722 [inlined]
   [3] walk(x::Expr, inner::Function, outer::ReactantCore.var"#28#29")
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:135
   [4] postwalk
     @ ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145 [inlined]
   [5] (::MacroTools.var"#25#26"{ReactantCore.var"#28#29"})(x::Expr)
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145
   [6] iterate
     @ ./generator.jl:48 [inlined]
   [7] collect_to!(dest::Vector{Any}, itr::Base.Generator{Vector{…}, MacroTools.var"#25#26"{…}}, offs::Int64, st::Int64)
     @ Base ./array.jl:849
   [8] collect_to!(dest::Vector{…}, itr::Base.Generator{…}, offs::Int64, st::Int64)
     @ Base ./array.jl:857
   [9] collect_to_with_first!(dest::Vector{…}, v1::LineNumberNode, itr::Base.Generator{…}, st::Int64)
     @ Base ./array.jl:827
  [10] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
     @ Base ./array.jl:821
  [11] collect_similar(cont::Vector{Any}, itr::Base.Generator{Vector{Any}, MacroTools.var"#25#26"{ReactantCore.var"#28#29"}})
     @ Base ./array.jl:720
  [12] map(f::Function, A::Vector{Any})
     @ Base ./abstractarray.jl:3371
  [13] walk(x::Expr, inner::Function, outer::ReactantCore.var"#28#29")
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:135
  [14] postwalk
     @ ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145 [inlined]
  [15] (::MacroTools.var"#25#26"{ReactantCore.var"#28#29"})(x::Expr)
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145
  [16] iterate
     @ ./generator.jl:48 [inlined]
  [17] collect_to!(dest::Vector{Expr}, itr::Base.Generator{Vector{…}, MacroTools.var"#25#26"{…}}, offs::Int64, st::Int64)
     @ Base ./array.jl:849
--- the above 9 lines are repeated 3 more times ---
  [45] collect_to!(dest::Vector{…}, itr::Base.Generator{…}, offs::Int64, st::Int64)
     @ Base ./array.jl:857
  [46] collect_to_with_first!(dest::Vector{…}, v1::LineNumberNode, itr::Base.Generator{…}, st::Int64)
     @ Base ./array.jl:827
  [47] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
     @ Base ./array.jl:821
  [48] collect_similar(cont::Vector{Any}, itr::Base.Generator{Vector{Any}, MacroTools.var"#25#26"{ReactantCore.var"#28#29"}})
     @ Base ./array.jl:720
  [49] map(f::Function, A::Vector{Any})
     @ Base ./abstractarray.jl:3371
  [50] walk(x::Expr, inner::Function, outer::ReactantCore.var"#28#29")
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:135
  [51] postwalk
     @ ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145 [inlined]
  [52] (::MacroTools.var"#25#26"{ReactantCore.var"#28#29"})(x::Expr)
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145
  [53] iterate
     @ ./generator.jl:48 [inlined]
  [54] collect_to!(dest::Vector{Expr}, itr::Base.Generator{Vector{…}, MacroTools.var"#25#26"{…}}, offs::Int64, st::Int64)
     @ Base ./array.jl:849
--- the above 9 lines are repeated 1 more time ---
--- the above 19 lines are repeated 1 more time ---
  [83] collect_to_with_first!(dest::Vector{…}, v1::LineNumberNode, itr::Base.Generator{…}, st::Int64)
     @ Base ./array.jl:827
  [84] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
     @ Base ./array.jl:821
  [85] collect_similar(cont::Vector{Any}, itr::Base.Generator{Vector{Any}, MacroTools.var"#25#26"{ReactantCore.var"#28#29"}})
     @ Base ./array.jl:720
  [86] map(f::Function, A::Vector{Any})
     @ Base ./abstractarray.jl:3371
  [87] walk(x::Expr, inner::Function, outer::ReactantCore.var"#28#29")
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:135
  [88] postwalk
     @ ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145 [inlined]
  [89] (::MacroTools.var"#25#26"{ReactantCore.var"#28#29"})(x::Expr)
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145
  [90] iterate
     @ ./generator.jl:48 [inlined]
  [91] collect_to!(dest::Vector{…}, itr::Base.Generator{…}, offs::Int64, st::Int64)
     @ Base ./array.jl:849
--- the above 9 lines are repeated 6 more times ---
 [146] collect_to!(dest::Vector{…}, itr::Base.Generator{…}, offs::Int64, st::Int64)
     @ Base ./array.jl:857
 [147] collect_to_with_first!(dest::Vector{…}, v1::LineNumberNode, itr::Base.Generator{…}, st::Int64)
     @ Base ./array.jl:827
 [148] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
     @ Base ./array.jl:821
 [149] collect_similar(cont::Vector{Any}, itr::Base.Generator{Vector{Any}, MacroTools.var"#25#26"{ReactantCore.var"#28#29"}})
     @ Base ./array.jl:720
 [150] map(f::Function, A::Vector{Any})
     @ Base ./abstractarray.jl:3371
 [151] walk(x::Expr, inner::Function, outer::ReactantCore.var"#28#29")
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:135
 [152] postwalk
     @ ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145 [inlined]
 [153] (::MacroTools.var"#25#26"{ReactantCore.var"#28#29"})(x::Expr)
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145
 [154] iterate
     @ ./generator.jl:48 [inlined]
 [155] collect_to!(dest::Vector{Any}, itr::Base.Generator{Vector{…}, MacroTools.var"#25#26"{…}}, offs::Int64, st::Int64)
     @ Base ./array.jl:849
 [156] collect_to!(dest::Vector{…}, itr::Base.Generator{…}, offs::Int64, st::Int64)
     @ Base ./array.jl:857
 [157] collect_to_with_first!(dest::Vector{…}, v1::LineNumberNode, itr::Base.Generator{…}, st::Int64)
     @ Base ./array.jl:827
 [158] _collect(c::Vector{…}, itr::Base.Generator{…}, ::Base.EltypeUnknown, isz::Base.HasShape{…})
     @ Base ./array.jl:821
 [159] collect_similar(cont::Vector{Any}, itr::Base.Generator{Vector{Any}, MacroTools.var"#25#26"{ReactantCore.var"#28#29"}})
     @ Base ./array.jl:720
 [160] map(f::Function, A::Vector{Any})
     @ Base ./abstractarray.jl:3371
 [161] walk(x::Expr, inner::Function, outer::ReactantCore.var"#28#29")
     @ MacroTools ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:135
 [162] postwalk
     @ ~/.julia/packages/MacroTools/Ar0jT/src/utils.jl:145 [inlined]
 [163] error_if_any_control_flow
     @ ~/.julia/packages/ReactantCore/SEqVX/src/ReactantCore.jl:719 [inlined]
 [164] trace_for(expr::Expr; track_numbers::Type, checkpointing::Bool, mincut::Bool)
     @ ReactantCore ~/.julia/packages/ReactantCore/SEqVX/src/ReactantCore.jl:383
in expression starting at /Users/guillaumedalle/Documents/GitHub/Julia/_Scratchpad/mwe.jl:18
Some type information was truncated. Use `show(err)` to see complete types.

[wsmoses]

oh yeah this is a known issue, where I think the julia expr rewriter in the traced if needs to be rewritten to look a bit more like the current impl within traced for.

Want to give it a go? It's all pure julia

[gdalle]

Where does that code live?

[wsmoses]

mostly

Reactant.jl/lib/ReactantCore/src/ReactantCore.jl

Line 615 in 1efd123

($(all_output_vars...),) = $(traced_if)(

@Pangoraw may be able to help give some pointers

[Pangoraw]

Essentially, the check that prohibits return, continue and such from within @trace should stop at function boundaries (instead of using MacroTools.postwalk).

Here is a minimal reproducer:

julia> ReactantCore.error_if_any_control_flow(:(
           begin
               map(f -> return f, arr)
           end
       ))
ERROR: Cannot use @trace on a block that contains a return statement

Reactant.jl/lib/ReactantCore/src/ReactantCore.jl

Lines 718 to 727 in 837e752

	function error_if_any_control_flow(expr)
	return MacroTools.postwalk(expr) do x
	for head in CONTROL_FLOW_EXPRS
	if Meta.isexpr(x, head)
	error("Cannot use @trace on a block that contains a $head statement")
	end
	end
	return x
	end
	end