Merge pull request #127 from SymbolicML/n-arity-v4

Permit n-argument operators

Author: MilesCranmer

.github/codecov.yml

@@ -0,0 +1,5 @@
+coverage:
+  status:
+    patch:
+      default:
+        informational: true

Project.toml

@@ -1,10 +1,11 @@
 name = "DynamicExpressions"
 uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 authors = ["MilesCranmer <[email protected]>"]
-version = "1.10.3"
+version = "2.0.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
 Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
@@ -30,6 +31,7 @@ DynamicExpressionsZygoteExt = "Zygote"
 [compat]
 Bumper = "0.6"
 ChainRulesCore = "1"
+Compat = "4.16"
 DispatchDoctor = "0.4"
 Interfaces = "0.3"
 LoopVectorization = "0.12"

README.md

@@ -27,7 +27,7 @@ A dynamic expression is a snippet of code that can change throughout runtime - c
 ```julia
 using DynamicExpressions
 
-operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos])
+operators = OperatorEnum(1 => (cos,), 2 => (+, -, *))
 variable_names = ["x1", "x2"]
 
 x1 = Expression(Node{Float64}(feature=1); operators, variable_names)
@@ -98,7 +98,7 @@ We can also compute gradients with the same speed:
 ```julia
 using Zygote  # trigger extension
 
-operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos])
+operators = OperatorEnum(1 => (cos,), 2 => (+, -, *))
 variable_names = ["x1", "x2"]
 x1, x2 = (Expression(Node{Float64}(feature=i); operators, variable_names) for i in 1:2)
 
@@ -149,7 +149,7 @@ using DynamicExpressions: @declare_expression_operator
 my_string_func(x::String) = "ello $x"
 @declare_expression_operator(my_string_func, 1)
 
-operators = GenericOperatorEnum(; binary_operators=[*], unary_operators=[my_string_func])
+operators = GenericOperatorEnum(1 => (my_string_func,), 2 => (*,))
 
 x1 = Expression(_x1; operators, variable_names)
 ```
@@ -192,7 +192,7 @@ vec_square(x) = x .* x
 @declare_expression_operator(vec_square, 1)
 
 # Set up an operator enum:
-operators = GenericOperatorEnum(;binary_operators=[vec_add], unary_operators=[vec_square])
+operators = GenericOperatorEnum(1 => (vec_square,), 2 => (vec_add,))
 
 # Construct the expression:
 variable_names = ["x1"]

docs/Project.toml

@@ -4,3 +4,6 @@ DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[sources]
+DynamicExpressions = { path = "../" }

docs/src/api.md

@@ -10,17 +10,11 @@ This `enum` is defined as follows:
 OperatorEnum
 ```
 
-Construct this operator specification as follows:
-
-```@docs
-OperatorEnum(; binary_operators=[], unary_operators=[], define_helper_functions::Bool=true)
-```
-
 This is just for scalar operators. However, you can use
 the following for more general operators:
 
 ```@docs
-GenericOperatorEnum(; binary_operators=[], unary_operators=[], define_helper_functions::Bool=true)
+GenericOperatorEnum
 ```
 
 By default, these operators will define helper functions for constructing trees,
@@ -60,7 +54,7 @@ When using these node constructors, types will automatically be promoted.
 You can convert the type of a node using `convert`:
 
 ```@docs
-convert(::Type{AbstractExpressionNode{T1}}, tree::AbstractExpressionNode{T2}) where {T1, T2}
+convert(::Type{N1}, tree::N2) where {T1,T2,D1,D2,N1<:AbstractExpressionNode{T1,D1},N2<:AbstractExpressionNode{T2,D2}}
 ```
 
 You can set a `tree` (in-place) with `set_node!`:
@@ -75,6 +69,41 @@ You can create a copy of a node with `copy_node`:
 copy_node
 ```
 
+## Generic Node Accessors
+
+For working with nodes of arbitrary arity:
+
+```@docs
+get_child
+set_child!
+get_children
+set_children!
+```
+
+Examples:
+
+```julia
+# Define operators including ternary
+my_ternary(x, y, z) = x + y * z
+operators = OperatorEnum(((sin,), (+, *), (my_ternary,)))  # (unary, binary, ternary)
+
+tree = Node{Float64,3}(; op=1, children=(Node{Float64,3}(; val=1.0), Node{Float64,3}(; val=2.0)))
+new_child = Node{Float64,3}(; val=3.0)
+
+left_child = get_child(tree, 1)
+right_child = get_child(tree, 2)
+
+set_child!(tree, new_child, 1)
+
+left, right = get_children(tree, Val(2))  # type stable
+
+# Transform to ternary operation
+child1, child2, child3 = Node{Float64,3}(; val=4.0), Node{Float64,3}(; val=5.0), Node{Float64,3}(; val=6.0)
+set_children!(tree, (child1, child2, child3))
+tree.op = 1  # my_ternary
+tree.degree = 3
+```
+
 ## Graph Nodes
 
 You can describe an equation as a *graph* rather than a tree
@@ -88,9 +117,7 @@ This makes it so you can have multiple parents for a given node,
 and share parts of an expression. For example:
 
 ```julia
-julia> operators = OperatorEnum(;
-           binary_operators=[+, -, *], unary_operators=[cos, sin, exp]
-       );
+julia> operators = OperatorEnum(1 => (cos, sin, exp), 2 => (+, -, *));
 
 julia> x1, x2 = GraphNode(feature=1), GraphNode(feature=2)
 (x1, x2)
@@ -109,7 +136,7 @@ This means that we only need to change it once
 to have changes propagate across the expression:
 
 ```julia
-julia> y.r.val *= 0.9
+julia> get_child(y, 2).val *= 0.9
 1.35
 
 julia> z

docs/src/eval.md

@@ -40,7 +40,7 @@ For example,
 ```@example
 using DynamicExpressions
 
-operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos])
+operators = OperatorEnum(1 => (cos, sin), 2 => (+, -, *, /))
 tree = Node(; feature=1) * cos(Node(; feature=2) - 3.2)
 
 tree([1 2 3; 4 5 6.], operators)
@@ -155,7 +155,7 @@ Let's look at an example. First, let's create a tree:
 ```julia
 using DynamicExpressions
 
-operators = OperatorEnum(binary_operators=(+, -, *, /), unary_operators=(cos, sin))
+operators = OperatorEnum(1 => (cos, sin), 2 => (+, -, *, /))
 
 x1 = Node{Float64}(feature=1)
 x2 = Node{Float64}(feature=2)

docs/src/utils.md

@@ -15,7 +15,6 @@ mapreduce(f::F, op::G, tree::AbstractNode; return_type, f_on_shared, break_shari
 any(f::F, tree::AbstractNode) where {F<:Function}
 all(f::F, tree::AbstractNode) where {F<:Function}
 map(f::F, tree::AbstractNode, result_type::Type{RT}=Nothing; break_sharing::Val=Val(false)) where {F<:Function,RT}
-convert(::Type{<:AbstractExpressionNode{T1}}, n::AbstractExpressionNode{T2}) where {T1,T2}
 hash(tree::AbstractExpressionNode{T}, h::UInt; break_sharing::Val=Val(false)) where {T}
 ```
 

ext/DynamicExpressionsBumperExt.jl

@@ -5,8 +5,7 @@ using DynamicExpressions:
     OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array, EvalOptions
 using DynamicExpressions.UtilsModule: ResultOk, counttuple
 
-import DynamicExpressions.ExtensionInterfaceModule:
-    bumper_eval_tree_array, bumper_kern1!, bumper_kern2!
+import DynamicExpressions.ExtensionInterfaceModule: bumper_eval_tree_array, bumper_kern!
 
 function bumper_eval_tree_array(
     tree::AbstractExpressionNode{T},
@@ -37,8 +36,7 @@ function bumper_eval_tree_array(
             branch_node -> branch_node,
             # In the evaluation kernel, we combine the branch nodes
             # with the arrays created by the leaf nodes:
-            ((args::Vararg{Any,M}) where {M}) ->
-                dispatch_kerns!(operators, args..., eval_options),
+            KernelDispatcher(operators, eval_options),
             tree;
             break_sharing=Val(true),
         )
@@ -49,63 +47,44 @@ function bumper_eval_tree_array(
     return (result, all_ok[])
 end
 
-function dispatch_kerns!(
-    operators, branch_node, cumulator, eval_options::EvalOptions{<:Any,true,early_exit}
-) where {early_exit}
-    cumulator.ok || return cumulator
-
-    out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, eval_options)
-    return ResultOk(out, early_exit ? is_valid_array(out) : true)
-end
-function dispatch_kerns!(
-    operators,
-    branch_node,
-    cumulator1,
-    cumulator2,
-    eval_options::EvalOptions{<:Any,true,early_exit},
-) where {early_exit}
-    cumulator1.ok || return cumulator1
-    cumulator2.ok || return cumulator2
-
-    out = dispatch_kern2!(
-        operators.binops, branch_node.op, cumulator1.x, cumulator2.x, eval_options
-    )
-    return ResultOk(out, early_exit ? is_valid_array(out) : true)
+struct KernelDispatcher{O<:OperatorEnum,E<:EvalOptions{<:Any,true,<:Any}} <: Function
+    operators::O
+    eval_options::E
 end
 
-@generated function dispatch_kern1!(unaops, op_idx, cumulator, eval_options::EvalOptions)
-    nuna = counttuple(unaops)
+@generated function (kd::KernelDispatcher{<:Any,<:EvalOptions{<:Any,true,early_exit}})(
+    branch_node, inputs::Vararg{Any,degree}
+) where {degree,early_exit}
     quote
-        Base.@nif(
-            $nuna,
-            i -> i == op_idx,
-            i -> let op = unaops[i]
-                return bumper_kern1!(op, cumulator, eval_options)
-            end,
-        )
+        Base.Cartesian.@nexprs($degree, i -> inputs[i].ok || return inputs[i])
+        cumulators = Base.Cartesian.@ntuple($degree, i -> inputs[i].x)
+        out = dispatch_kerns!(kd.operators, branch_node, cumulators, kd.eval_options)
+        return ResultOk(out, early_exit ? is_valid_array(out) : true)
     end
 end
-@generated function dispatch_kern2!(
-    binops, op_idx, cumulator1, cumulator2, eval_options::EvalOptions
-)
-    nbin = counttuple(binops)
+@generated function dispatch_kerns!(
+    operators::OperatorEnum{OPS},
+    branch_node,
+    cumulators::Tuple{Vararg{Any,degree}},
+    eval_options::EvalOptions,
+) where {OPS,degree}
+    nops = length(OPS.types[degree].types)
     quote
-        Base.@nif(
-            $nbin,
+        op_idx = branch_node.op
+        Base.Cartesian.@nif(
+            $nops,
             i -> i == op_idx,
-            i -> let op = binops[i]
-                return bumper_kern2!(op, cumulator1, cumulator2, eval_options)
-            end,
+            i -> bumper_kern!(operators[$degree][i], cumulators, eval_options)
         )
     end
 end
-function bumper_kern1!(op::F, cumulator, ::EvalOptions{false,true}) where {F}
-    @. cumulator = op(cumulator)
-    return cumulator
-end
-function bumper_kern2!(op::F, cumulator1, cumulator2, ::EvalOptions{false,true}) where {F}
-    @. cumulator1 = op(cumulator1, cumulator2)
-    return cumulator1
+
+function bumper_kern!(
+    op::F, cumulators::Tuple{Vararg{Any,degree}}, ::EvalOptions{false,true,early_exit}
+) where {F,degree,early_exit}
+    cumulator_1 = first(cumulators)
+    @. cumulator_1 = op(cumulators...)
+    return cumulator_1
 end
 
 end