diff --git a/ext/DynamicExpressionsBumperExt.jl b/ext/DynamicExpressionsBumperExt.jl index 31b234bb..a7b40e16 100644 --- a/ext/DynamicExpressionsBumperExt.jl +++ b/ext/DynamicExpressionsBumperExt.jl @@ -5,9 +5,10 @@ using DynamicExpressions: OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array, EvalContext using DynamicExpressions.UtilsModule: ResultOk, counttuple -import DynamicExpressions.ExtensionInterfaceModule: bumper_eval_tree_array, bumper_kern! +import DynamicExpressions.ExtensionInterfaceModule: + _bumper_eval_tree_array, is_extension_loaded, bumper_kern! -function bumper_eval_tree_array( +function _bumper_eval_tree_array( tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, operators::OperatorEnum, @@ -87,4 +88,6 @@ function bumper_kern!( return cumulator_1 end +is_extension_loaded(::Val{:Bumper}) = true + end diff --git a/ext/DynamicExpressionsSymbolicUtilsExt.jl b/ext/DynamicExpressionsSymbolicUtilsExt.jl index 637faa94..4b419403 100644 --- a/ext/DynamicExpressionsSymbolicUtilsExt.jl +++ b/ext/DynamicExpressionsSymbolicUtilsExt.jl @@ -10,7 +10,12 @@ using DynamicExpressions.UtilsModule: deprecate_varmap using SymbolicUtils using SymbolicUtils: BasicSymbolic, SymReal, iscall, issym, isconst, unwrap_const -import DynamicExpressions.ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node +import DynamicExpressions.ExtensionInterfaceModule: + is_extension_loaded, + _node_to_symbolic, + _symbolic_to_node, + node_to_symbolic, + symbolic_to_node import DynamicExpressions.ValueInterfaceModule: is_valid const SYMBOLIC_UTILS_TYPES = Union{<:Number,BasicSymbolic} @@ -268,25 +273,7 @@ function Base.convert( return constructorof(E)(tree; operators, variable_names, kws...) end -""" - node_to_symbolic(tree::AbstractExpressionNode, operators::AbstractOperatorEnum; - variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing, - index_functions::Bool=false) - -The interface to SymbolicUtils.jl. Passing a tree to this function -will generate a symbolic equation in SymbolicUtils.jl format. - -## Arguments - -- `tree::AbstractExpressionNode`: The equation to convert. -- `operators::AbstractOperatorEnum`: OperatorEnum, which contains the operators used in the equation. -- `variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing`: What variable names to use for - each feature. Default is [x1, x2, x3, ...]. -- `index_functions::Bool=false`: Whether to generate special names for the - operators, which then allows one to convert back to a `AbstractExpressionNode` format - using `symbolic_to_node`. -""" -function node_to_symbolic( +function _node_to_symbolic( tree::AbstractExpressionNode{T,2}, operators::AbstractOperatorEnum; variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing, @@ -309,7 +296,7 @@ function node_to_symbolic( ) return substitute(expr, subs) end -function node_to_symbolic( +function _node_to_symbolic( tree::AbstractExpression, operators::Union{AbstractOperatorEnum,Nothing}=nothing; variable_names::Union{AbstractVector{<:AbstractString},Nothing}=nothing, @@ -323,7 +310,7 @@ function node_to_symbolic( ) end -function symbolic_to_node( +function _symbolic_to_node( eqn::BasicSymbolic, operators::AbstractOperatorEnum, ::Type{N}=Node; @@ -429,4 +416,6 @@ function multiply_powers( end end +is_extension_loaded(::Val{:SymbolicUtils}) = true + end diff --git a/ext/DynamicExpressionsZygoteExt.jl b/ext/DynamicExpressionsZygoteExt.jl index f42a3f89..57346821 100644 --- a/ext/DynamicExpressionsZygoteExt.jl +++ b/ext/DynamicExpressionsZygoteExt.jl @@ -1,9 +1,10 @@ module DynamicExpressionsZygoteExt using Zygote: gradient -import DynamicExpressions.ExtensionInterfaceModule: _zygote_gradient, ZygoteGradient +import DynamicExpressions.ExtensionInterfaceModule: + is_extension_loaded, _zygote_gradient_impl, ZygoteGradient -function _zygote_gradient(op::F, ::Val{degree}) where {F,degree} +function _zygote_gradient_impl(op::F, ::Val{degree}) where {F,degree} return ZygoteGradient{F,degree}(op) end @@ -13,4 +14,6 @@ function (g::ZygoteGradient{F,degree})(args::Vararg{Any,degree}) where {F,degree return ntuple(i -> @something(partials[i], zero(args[i])), Val(degree)) end +is_extension_loaded(::Val{:Zygote}) = true + end diff --git a/src/ExtensionInterface.jl b/src/ExtensionInterface.jl index ca42430f..3ebce3eb 100644 --- a/src/ExtensionInterface.jl +++ b/src/ExtensionInterface.jl @@ -1,11 +1,40 @@ module ExtensionInterfaceModule -function node_to_symbolic(args...; kws...) - return error("Please load the `SymbolicUtils` package to use `node_to_symbolic`.") +using DispatchDoctor: @unstable + +is_extension_loaded(::Val) = false + +""" + node_to_symbolic(tree::AbstractExpressionNode, operators::AbstractOperatorEnum; + variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing, + index_functions::Bool=false) + +The interface to SymbolicUtils.jl. Passing a tree to this function +will generate a symbolic equation in SymbolicUtils.jl format. + +## Arguments + +- `tree::AbstractExpressionNode`: The equation to convert. +- `operators::AbstractOperatorEnum`: OperatorEnum, which contains the operators used in the equation. +- `variable_names::Union{AbstractVector{<:AbstractString}, Nothing}=nothing`: What variable names to use for + each feature. Default is [x1, x2, x3, ...]. +- `index_functions::Bool=false`: Whether to generate special names for the + operators, which then allows one to convert back to a `AbstractExpressionNode` format + using `symbolic_to_node`. +""" +@unstable function node_to_symbolic(args...; kws...) + is_extension_loaded(Val(:SymbolicUtils)) || + error("Please load the `SymbolicUtils` package to use `node_to_symbolic`.") + return _node_to_symbolic(args...; kws...) end -function symbolic_to_node(args...; kws...) - return error("Please load the `SymbolicUtils` package to use `symbolic_to_node`.") +function _node_to_symbolic end + +@unstable function symbolic_to_node(args...; kws...) + is_extension_loaded(Val(:SymbolicUtils)) || + error("Please load the `SymbolicUtils` package to use `symbolic_to_node`.") + return _symbolic_to_node(args...; kws...) end +function _symbolic_to_node end struct ZygoteGradient{F,degree} <: Function op::F @@ -19,12 +48,17 @@ end Base.show(io::IO, ::MIME"text/plain", g::ZygoteGradient) = show(io, g) function _zygote_gradient(args...) - return error("Please load the Zygote.jl package.") + is_extension_loaded(Val(:Zygote)) || error("Please load the Zygote.jl package.") + return _zygote_gradient_impl(args...) end +function _zygote_gradient_impl end function bumper_eval_tree_array(args...) - return error("Please load the Bumper.jl package to use this feature.") + is_extension_loaded(Val(:Bumper)) || + error("Please load the Bumper.jl package to use this feature.") + return _bumper_eval_tree_array(args...) end +function _bumper_eval_tree_array end function bumper_kern! end _is_loopvectorization_loaded(_) = false # COV_EXCL_LINE diff --git a/test/test_initial_errors.jl b/test/test_initial_errors.jl index 5a9bf01e..fad1cc39 100644 --- a/test/test_initial_errors.jl +++ b/test/test_initial_errors.jl @@ -1,5 +1,7 @@ using DynamicExpressions using DynamicExpressions: EvalContext +using DynamicExpressions.ExtensionInterfaceModule: + is_extension_loaded, _zygote_gradient, bumper_eval_tree_array using DispatchDoctor: allow_unstable using Test @@ -46,3 +48,41 @@ tree = cos(2.1 * x1) + sin(x2) () -> tree(ones(2, 10), operators; eval_context=EvalContext(; turbo=Val(true))) ) ) + +# Loaded extensions should use normal dispatch for unsupported arguments instead of claiming +# that the dependency is missing. +using SymbolicUtils + +@test is_extension_loaded(Val(:SymbolicUtils)) +@test !is_extension_loaded(Val(:Zygote)) +@test !is_extension_loaded(Val(:Bumper)) +@test_throws "Please load the Zygote.jl package." _zygote_gradient(nothing) +@test_throws "Please load the Bumper.jl package" bumper_eval_tree_array(nothing) + +symbolic_x1 = allow_unstable(() -> node_to_symbolic(x1, operators)) +@test string(symbolic_x1) == "x1" +@test string(allow_unstable(() -> symbolic_to_node(symbolic_x1, operators))) == "x1" +@test_throws MethodError node_to_symbolic(nothing) +@test_throws MethodError symbolic_to_node(nothing) + +using Zygote + +@test is_extension_loaded(Val(:SymbolicUtils)) +@test is_extension_loaded(Val(:Zygote)) +@test !is_extension_loaded(Val(:Bumper)) +@test_throws "Please load the Bumper.jl package" bumper_eval_tree_array(nothing) + +@test allow_unstable(() -> only(_zygote_gradient(sin, Val(1))(1.0))) ≈ cos(1.0) +@test_throws MethodError _zygote_gradient(nothing) + +using Bumper + +@test is_extension_loaded(Val(:SymbolicUtils)) +@test is_extension_loaded(Val(:Zygote)) +@test is_extension_loaded(Val(:Bumper)) + +bumper_result = allow_unstable( + () -> tree(ones(2, 10), operators; eval_options=EvalContext(; bumper=Val(true))) +) +@test bumper_result ≈ fill(cos(2.1) + sin(1.0), 10) +@test_throws MethodError bumper_eval_tree_array(nothing)