Merge pull request #79 from SymbolicML/dev

Add DispatchDoctor.jl and fix various type instabilities

Author: MilesCranmer

.github/workflows/CI.yml

@@ -61,8 +61,8 @@ jobs:
           path-to-lcov: lcov.info
           flag-name: julia-${{ matrix.julia-version }}-${{ matrix.os }}-main-${{ github.event_name }}
 
-  integration_tests:
-    name: Integration test - ${{ matrix.test_name }} - ${{ matrix.os }}
+  additional_tests:
+    name: test ${{ matrix.test_name }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 60
     strategy:
@@ -74,6 +74,7 @@ jobs:
           - "1"
         test_name:
           - "enzyme"
+          - "jet"
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
@@ -100,7 +101,7 @@ jobs:
     runs-on: ubuntu-latest
     needs:
       - test
-      - integration_tests
+      - additional_tests
     steps:
       - name: Finish
         uses: coverallsapp/github-action@v2

Project.toml

@@ -1,11 +1,12 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <miles.cranmer@gmail.com>"]
-version = "0.18.0-alpha"
+version = "0.18.0-alpha.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
@@ -31,6 +32,7 @@ DynamicExpressionsZygoteExt = "Zygote"
 Bumper = "0.6"
 ChainRulesCore = "1"
 Compat = "3.37, 4"
+DispatchDoctor = "0.4"
 LoopVectorization = "0.12"
 MacroTools = "0.4, 0.5"
 Optim = "0.19, 1"

src/ChainRules.jl

@@ -38,10 +38,10 @@ function ChainRulesCore.rrule(
         dtree = let X = X, dY = dY, tree = tree, operators = operators
             @thunk(
                 let
-                    _, gradient, complete = eval_grad_tree_array(
+                    _, gradient, complete2 = eval_grad_tree_array(
                         tree, X, operators; variable=Val(false)
                     )
-                    if !complete
+                    if !complete2
                         gradient .= NaN
                     end
 
@@ -55,14 +55,14 @@ function ChainRulesCore.rrule(
         dX = let X = X, dY = dY, tree = tree, operators = operators
             @thunk(
                 let
-                    _, gradient, complete = eval_grad_tree_array(
+                    _, gradient2, complete3 = eval_grad_tree_array(
                         tree, X, operators; variable=Val(true)
                     )
-                    if !complete
-                        gradient .= NaN
+                    if !complete3
+                        gradient2 .= NaN
                     end
 
-                    gradient .* reshape(dY, 1, length(dY))
+                    gradient2 .* reshape(dY, 1, length(dY))
                 end
             )
         end

src/DynamicExpressions.jl

@@ -1,23 +1,29 @@
 module DynamicExpressions
 
-include("Utils.jl")
-include("ExtensionInterface.jl")
-include("OperatorEnum.jl")
-include("Node.jl")
-include("NodeUtils.jl")
-include("Strings.jl")
-include("Evaluate.jl")
-include("EvaluateDerivative.jl")
-include("ChainRules.jl")
-include("EvaluationHelpers.jl")
-include("Simplify.jl")
-include("OperatorEnumConstruction.jl")
-include("Random.jl")
-include("Expression.jl")
-include("Parse.jl")
+using DispatchDoctor: @stable, @unstable
+
+@stable default_mode = "disable" begin
+    include("Utils.jl")
+    include("ExtensionInterface.jl")
+    include("OperatorEnum.jl")
+    include("Node.jl")
+    include("NodeUtils.jl")
+    include("Strings.jl")
+    include("Evaluate.jl")
+    include("EvaluateDerivative.jl")
+    include("ChainRules.jl")
+    include("EvaluationHelpers.jl")
+    include("Simplify.jl")
+    include("OperatorEnumConstruction.jl")
+    include("Random.jl")
+    include("Expression.jl")
+    include("Parse.jl")
+end
 
 import PackageExtensionCompat: @require_extensions
 import Reexport: @reexport
+macro ignore(args...) end
+
 @reexport import .NodeModule:
     AbstractNode,
     AbstractExpressionNode,
@@ -86,11 +92,9 @@ const PACKAGE_VERSION = let d = pkgdir(@__MODULE__)
     end
 end
 
-macro ignore(args...) end
 # To get LanguageServer to register library within tests
 @ignore include("../test/runtests.jl")
 
 include("precompile.jl")
 do_precompilation(; mode=:precompile)
-
 end

src/Evaluate.jl

@@ -1,5 +1,7 @@
 module EvaluateModule
 
+using DispatchDoctor: @unstable
+
 import ..NodeModule: AbstractExpressionNode, constructorof
 import ..StringsModule: string_tree
 import ..OperatorEnumModule: OperatorEnum, GenericOperatorEnum
@@ -657,7 +659,7 @@ function eval(current_node)
     A `false` complete means an operator was called on input types
     that it was not defined for.
 """
-function eval_tree_array(
+@unstable function eval_tree_array(
     tree::AbstractExpressionNode,
     cX::AbstractArray,
     operators::GenericOperatorEnum;
@@ -680,7 +682,7 @@ function eval_tree_array(
     end
 end
 
-function _eval_tree_array_generic(
+@unstable function _eval_tree_array_generic(
     tree::AbstractExpressionNode{T1},
     cX::AbstractArray{T2,N},
     operators::GenericOperatorEnum,
@@ -707,7 +709,7 @@ function _eval_tree_array_generic(
     end
 end
 
-function deg1_eval_generic(
+@unstable function deg1_eval_generic(
     tree, cX, op::F, operators::GenericOperatorEnum, ::Val{throw_errors}
 ) where {F,throw_errors}
     left, complete = eval_tree_array(tree.l, cX, operators)
@@ -716,7 +718,7 @@ function deg1_eval_generic(
     return op(left), true
 end
 
-function deg2_eval_generic(
+@unstable function deg2_eval_generic(
     tree, cX, op::F, operators::GenericOperatorEnum, ::Val{throw_errors}
 ) where {F,throw_errors}
     left, complete = eval_tree_array(tree.l, cX, operators)

src/Expression.jl

@@ -1,6 +1,7 @@
 """This module defines a user-facing `Expression` type"""
 module ExpressionModule
 
+using DispatchDoctor: @unstable
 using ..NodeModule: AbstractExpressionNode
 using ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum
 using ..UtilsModule: Undefined
@@ -12,7 +13,7 @@ end
 _data(x::Metadata) = getfield(x, :_data)
 
 Base.propertynames(x::Metadata) = propertynames(_data(x))
-@inline Base.getproperty(x::Metadata, f::Symbol) = getproperty(_data(x), f)
+@unstable @inline Base.getproperty(x::Metadata, f::Symbol) = getproperty(_data(x), f)
 Base.show(io::IO, x::Metadata) = print(io, "Metadata(", _data(x), ")")
 @inline _copy(x) = copy(x)
 @inline _copy(x::Nothing) = nothing

src/OperatorEnumConstruction.jl

@@ -1,5 +1,7 @@
 module OperatorEnumConstructionModule
 
+using DispatchDoctor: @unstable
+
 import ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum, GenericOperatorEnum
 import ..NodeModule: Node, GraphNode, AbstractExpressionNode, constructorof
 import ..StringsModule: string_tree
@@ -43,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,
@@ -62,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,
@@ -93,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(
@@ -364,7 +366,7 @@ redefine operators for `AbstractExpressionNode` types, as well as `show`, `print
    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 OperatorEnum(;
+@unstable function OperatorEnum(;
     binary_operators=Function[],
     unary_operators=Function[],
     define_helper_functions::Bool=true,
@@ -417,7 +419,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/Parse.jl

@@ -1,5 +1,7 @@
 module ParseModule
 
+using DispatchDoctor: @unstable
+
 using ..NodeModule: AbstractExpressionNode, Node, constructorof
 using ..OperatorEnumModule: AbstractOperatorEnum
 using ..OperatorEnumConstructionModule: OperatorEnum, empty_all_globals!
@@ -97,7 +99,7 @@ macro parse_expression(ex, kws...)
     )
 end
 
-function _parse_kws(kws)
+@unstable function _parse_kws(kws)
     # Initialize default values for operators and variable_names
     operators = nothing
     variable_names = nothing
@@ -189,7 +191,7 @@ function _parse_kws(kws)
 end
 
 """Parse an expression Julia `Expr` object."""
-function parse_expression(
+@unstable function parse_expression(
     ex;
     operators::AbstractOperatorEnum,
     variable_names::Union{AbstractVector,Nothing}=nothing,
@@ -215,7 +217,7 @@ end
 """An empty module for evaluation without collisions."""
 module EmptyModule end
 
-function _parse_expression(
+@unstable function _parse_expression(
     ex::Expr,
     operators::AbstractOperatorEnum,
     variable_names::Union{AbstractVector{<:AbstractString},Nothing},
@@ -242,7 +244,7 @@ function _parse_expression(
         func, args, operators, variable_names, N, E, evaluate_on; kws...
     )
 end
-function _parse_expression(
+@unstable function _parse_expression(
     func::F,
     args,
     operators::AbstractOperatorEnum,
@@ -331,7 +333,7 @@ function _parse_expression(
         )
     end
 end
-function _parse_expression(
+@unstable function _parse_expression(
     ex,
     operators::AbstractOperatorEnum,
     variable_names::Union{AbstractVector{<:AbstractString},Nothing},
@@ -343,7 +345,7 @@ function _parse_expression(
     return parse_leaf(ex, variable_names, node_type, expression_type; kws...)
 end
 
-function parse_leaf(
+@unstable function parse_leaf(
     ex,
     variable_names,
     node_type::Type{<:AbstractExpressionNode},

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/Strings.jl

@@ -14,17 +14,19 @@ const OP_NAMES = Base.ImmutableDict(
     "safe_pow" => "^",
 )
 
-function dispatch_op_name(::Val{2}, ::Nothing, idx)::Vector{Char}
-    return vcat(collect("binary_operator["), collect(string(idx)), [']'])
-end
-function dispatch_op_name(::Val{1}, ::Nothing, idx)::Vector{Char}
-    return vcat(collect("unary_operator["), collect(string(idx)), [']'])
-end
-function dispatch_op_name(::Val{2}, operators::AbstractOperatorEnum, idx)::Vector{Char}
-    return get_op_name(operators.binops[idx])
+function dispatch_op_name(::Val{deg}, ::Nothing, idx)::Vector{Char} where {deg}
+    if deg == 1
+        return vcat(collect("unary_operator["), collect(string(idx)), [']'])
+    else
+        return vcat(collect("binary_operator["), collect(string(idx)), [']'])
+    end
 end
-function dispatch_op_name(::Val{1}, operators::AbstractOperatorEnum, idx)::Vector{Char}
-    return get_op_name(operators.unaops[idx])
+function dispatch_op_name(::Val{deg}, operators::AbstractOperatorEnum, idx) where {deg}
+    if deg == 1
+        return get_op_name(operators.unaops[idx])::Vector{Char}
+    else
+        return get_op_name(operators.binops[idx])::Vector{Char}
+    end
 end
 
 @generated function get_op_name(op::F)::Vector{Char} where {F}
@@ -137,15 +139,22 @@ function string_tree(
 )::String where {T,F1<:Function,F2<:Function}
     variable_names = deprecate_varmap(variable_names, varMap, :string_tree)
     raw_output = tree_mapreduce(
-        leaf -> if leaf.constant
-            collect(f_constant(leaf.val))
-        else
-            collect(f_variable(leaf.feature, variable_names))
+        let f_constant = f_constant,
+            f_variable = f_variable,
+            variable_names = variable_names
+
+            (leaf,) -> if leaf.constant
+                collect(f_constant(leaf.val))::Vector{Char}
+            else
+                collect(f_variable(leaf.feature, variable_names))::Vector{Char}
+            end
         end,
-        branch -> if branch.degree == 1
-            dispatch_op_name(Val(1), operators, branch.op)
-        else
-            dispatch_op_name(Val(2), operators, branch.op)
+        let operators = operators
+            (branch,) -> if branch.degree == 1
+                dispatch_op_name(Val(1), operators, branch.op)::Vector{Char}
+            else
+                dispatch_op_name(Val(2), operators, branch.op)::Vector{Char}
+            end
         end,
         combine_op_with_inputs,
         tree,