fix: some instabilities in base

Author: MilesCranmer

src/DynamicExpressions.jl

@@ -2,7 +2,7 @@ module DynamicExpressions
 
 using DispatchDoctor: @stable, @unstable
 
-@stable default_mode="disable" begin
+@stable default_mode = "disable" begin
     include("Utils.jl")
     include("ExtensionInterface.jl")
     include("OperatorEnum.jl")

src/OperatorEnumConstruction.jl

@@ -45,7 +45,7 @@ function Base.show(io::IO, tree::AbstractExpressionNode)
         return print(io, string_tree(tree, latest_operators; kwargs...))
     end
 end
-function (tree::AbstractExpressionNode)(X; kws...)
+@unstable function (tree::AbstractExpressionNode)(X; kws...)
     Base.depwarn(
         "The `tree(X; kws...)` syntax is deprecated. Use `tree(X, operators; kws...)` instead.",
         :AbstractExpressionNode,
@@ -64,7 +64,7 @@ function (tree::AbstractExpressionNode)(X; kws...)
     end
 end
 
-function _grad_evaluator(tree::AbstractExpressionNode, X; kws...)
+@unstable function _grad_evaluator(tree::AbstractExpressionNode, X; kws...)
     Base.depwarn(
         "The `tree'(X; kws...)` syntax is deprecated. Use `tree'(X, operators; kws...)` instead.",
         :AbstractExpressionNode,
@@ -95,7 +95,7 @@ function set_default_operators!(operators::GenericOperatorEnum)
     return LATEST_OPERATORS_TYPE.x = IsGenericOperatorEnum
 end
 
-function lookup_op(@nospecialize(f), ::Val{degree}) where {degree}
+@unstable function lookup_op(@nospecialize(f), ::Val{degree}) where {degree}
     mapping = degree == 1 ? LATEST_UNARY_OPERATOR_MAPPING : LATEST_BINARY_OPERATOR_MAPPING
     if !haskey(mapping, f)
         error(
@@ -420,7 +420,7 @@ and `(::AbstractExpressionNode)(X)`.
    are *not* needed for the package to work; they are purely for convenience.
 - `empty_old_operators::Bool=true`: Whether to clear the old operators.
 """
-function GenericOperatorEnum(;
+@unstable function GenericOperatorEnum(;
     binary_operators=Function[],
     unary_operators=Function[],
     define_helper_functions::Bool=true,

src/Simplify.jl

@@ -122,13 +122,9 @@ end
 
 # Simplify tree
 function simplify_tree!(tree::AbstractExpressionNode, operators::AbstractOperatorEnum)
-    tree = tree_mapreduce(
-        identity,
-        (p, c...) -> combine_children!(operators, p, c...),
-        tree,
-        constructorof(typeof(tree));
+    return tree_mapreduce(
+        identity, (p, c...) -> combine_children!(operators, p, c...), tree, typeof(tree);
     )
-    return tree
 end
 
 end

src/base.jl

@@ -33,6 +33,7 @@ using ..UtilsModule: @memoize_on, @with_memoize, Undefined
         [f_branch::Function,]
         op::Function,
         tree::AbstractNode,
+        [result_type::Type=Undefined];
         f_on_shared::Function=(result, is_shared) -> result,
         break_sharing::Val=Val(false),
     )
@@ -78,11 +79,11 @@ function tree_mapreduce(
     f::F,
     op::G,
     tree::AbstractNode,
-    result_type::Type=Undefined;
+    result_type::Type{RT}=Undefined;
     f_on_shared::H=(result, is_shared) -> result,
     break_sharing=Val(false),
-) where {F<:Function,G<:Function,H<:Function}
-    return tree_mapreduce(f, f, op, tree, result_type; f_on_shared, break_sharing)
+) where {RT,F<:Function,G<:Function,H<:Function}
+    return tree_mapreduce(f, f, op, tree, RT; f_on_shared, break_sharing)
 end
 function tree_mapreduce(
     f_leaf::F1,
@@ -341,7 +342,7 @@ function count(
 end
 
 """
-    sum(f::Function, tree::AbstractNode; init=0, return_type=Undefined, f_on_shared=_default_shared_aggregation, break_sharing::Val=Val(false)) where {F<:Function}
+    sum(f::Function, tree::AbstractNode; return_type=Undefined, f_on_shared=_default_shared_aggregation, break_sharing::Val=Val(false)) where {F<:Function}
 
 Sum the results of a function over a tree. For graphs with shared nodes
 such as `GraphNode`, the function `f_on_shared` is called on the result
@@ -351,15 +352,14 @@ behavior).
 function sum(
     f::F,
     tree::AbstractNode;
-    init=0,
     return_type=Undefined,
     f_on_shared=_default_shared_aggregation,
     break_sharing::Val=Val(false),
 ) where {F<:Function}
     if preserve_sharing(typeof(tree))
         @assert typeof(return_type) !== Undefined "Must specify `return_type` as a keyword argument to `sum` if `preserve_sharing` is true."
     end
-    return tree_mapreduce(f, +, tree, return_type; f_on_shared, break_sharing) + init
+    return tree_mapreduce(f, +, tree, return_type; f_on_shared, break_sharing)
 end
 function _default_shared_aggregation(c, is_shared)
     return is_shared ? (false * c) : c
@@ -382,15 +382,15 @@ function mapreduce(
     f::F,
     op::G,
     tree::AbstractNode;
-    return_type=Undefined,
+    return_type::Type{T}=Undefined,
     f_on_shared=(c, is_shared) -> is_shared ? (false * c) : c,
     break_sharing::Val=Val(false),
-) where {F<:Function,G<:Function}
-    if preserve_sharing(typeof(tree))
-        @assert typeof(return_type) !== Undefined "Must specify `return_type` as a keyword argument to `mapreduce` if `preserve_sharing` is true."
+) where {T,F<:Function,G<:Function}
+    if preserve_sharing(typeof(tree)) && break_sharing === Val(false)
+        @assert T !== Undefined "Must specify `return_type` as a keyword argument to `mapreduce` if `preserve_sharing` is true."
     end
     return tree_mapreduce(
-        f, (n...) -> reduce(op, n), tree, return_type; f_on_shared, break_sharing
+        f, (p, c...) -> reduce(op, (p, c...)), tree, T; f_on_shared, break_sharing
     )
 end
 

test/test_deprecations.jl

@@ -2,6 +2,7 @@ using DynamicExpressions
 using Test
 using Zygote
 using Suppressor: @capture_err
+using DispatchDoctor: allow_unstable
 
 operators = OperatorEnum(; binary_operators=[+, -, *, /], unary_operators=[cos, sin])
 x1, x2 = Node{Float64}(; feature=1), Node{Float64}(; feature=2)
@@ -16,10 +17,11 @@ for constructor in (OperatorEnum, GenericOperatorEnum)
 
     constructor == GenericOperatorEnum && continue
 
-    VERSION >= v"1.9" &&
-        @test_logs (:warn, r"The `tree'\(X; kws...\)` syntax is deprecated.*") tree'(
-            [1.0; 2.0;;]
-        )
+    if VERSION >= v"1.9"
+        @test_logs (:warn, r"The `tree'\(X; kws...\)` syntax is deprecated.*") allow_unstable() do
+            tree'([1.0; 2.0;;])
+        end
+    end
 end
 
 if VERSION >= v"1.9"

test/test_initial_errors.jl

@@ -33,14 +33,17 @@ if VERSION >= v"1.9"
         symbolic_to_node(tree, operators)
     )
 
-    @test_throws("Please load the Zygote.jl package.", tree'(ones(2, 10)))
+    @test_throws(
+        "Please load the Zygote.jl package.", allow_unstable(() -> tree'(ones(2, 10)))
+    )
 
     @test_throws(
-        "Please load the Bumper.jl package", tree(ones(2, 10), operators; bumper=Val(true))
+        "Please load the Bumper.jl package",
+        allow_unstable(() -> tree(ones(2, 10), operators; bumper=Val(true)))
     )
 
     @test_throws(
         "Please load the LoopVectorization.jl package",
-        tree(ones(2, 10), operators; turbo=Val(true))
+        allow_unstable(() -> tree(ones(2, 10), operators; turbo=Val(true)))
     )
 end