Need to work with a large space of symbolic expressions which are unknown at runtime, but numerically evaluate them just as quickly as if they were hard-coded?
Then you need DynamicExpressions.jl!
DynamicExpressions.jl is the newly-exposed backbone of the symbolic regression libraries SymbolicRegression.jl and PySR. These are machine learning algorithms which discover equations which optimize an objective, rather than tuning parameters in a flexible pre-defined model. Refactoring the underlying data structure into its own minimal library will make it easier for others to use and extend.
A dynamic expression is a snippet of code that changes throughout runtime - compilation is not possible! DynamicExpressions.jl does the following:
- Defines an enum over user-specified operators.
- Using this enum, it defines a very lightweight and type-stable data structure for arbitrary expressions.
- It then generates specialized evaluation kernels for the space of potential operators.
- It also generates kernels for the first-order derivatives, using Zygote.jl.
- It can also operate on arbitrary other types (vectors, tensors, symbols, strings, or even unions) - see last part below.
It also has import and export functionality with SymbolicUtils.jl, so you can move your runtime expression into a CAS!
using DynamicExpressions
operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos])
x1 = Node(; feature=1)
x2 = Node(; feature=2)
expression = x1 * cos(x2 - 3.2)
X = randn(Float64, 2, 100);
expression(X) # 100-element Vector{Float64}
(We can construct this expression with normal operators, since calling OperatorEnum() will @eval new functions on Node that use the specified enum.)
Speed
First, what happens if we naively use Julia symbols to define and then evaluate this expression?
@btime eval(:(X[1, :] .* cos.(X[2, :] .- 3.2)))
# 117,000 ns
This is quite slow, meaning it will be hard to quickly search over the space of expressions. Letâs see how DynamicExpressions.jl compares:
@btime expression(X)
# 693 ns
Much faster! And we didnât even need to compile it. (Internally, this is calling eval_tree_array(expression, X, operators) - where operators has been pre-defined when we called OperatorEnum()).
If we change expression dynamically with a random number generator, it will have the same performance:
@btime begin
expression.op = rand(1:3) # random operator in [+, -, *]
expression(X)
end
# 842 ns
Now, letâs see the performance if we had hard-coded these expressions:
f(X) = X[1, :] .* cos.(X[2, :] .- 3.2)
@btime f(X)
# 708 ns
So, our dynamic expression evaluation is about the same (or even a bit faster) as evaluating a basic hard-coded expression! Letâs see if we can optimize the speed of the hard-coded version:
f_optimized(X) = begin
y = Vector{Float64}(undef, 100)
@inbounds @simd for i=1:100
y[i] = X[1, i] * cos(X[2, i] - 3.2)
end
y
end
@btime f_optimized(X)
# 526 ns
The DynamicExpressions.jl version is only 25% slower than one which has been optimized by hand into a single SIMD kernel! Not bad at all.
More importantly: we can change expression throughout runtime, and expect the same performance. This makes this data structure ideal for symbolic regression and other evaluation-based searches over expression trees.
Derivatives
We can also compute gradients with the same speed:
operators = OperatorEnum(;
binary_operators=[+, -, *],
unary_operators=[cos],
enable_autodiff=true,
)
x1 = Node(; feature=1)
x2 = Node(; feature=2)
expression = x1 * cos(x2 - 3.2)
We can take the gradient with respect to inputs with simply the ' character:
grad = expression'(X)
This is quite fast:
@btime expression'(X)
# 2894 ns
and again, we can change this expression at runtime, without loss in performance!
@btime begin
expression.op = rand(1:3)
expression'(X)
end
# 3198 ns
Internally, this is calling the eval_grad_tree_array function, which performs forward-mode automatic differentiation on the expression tree with Zygote-compiled kernels. We can also compute the derivative with respect to constants:
result, grad, did_finish = eval_grad_tree_array(expression, X, operators; variable=false)
or with respect to variables, and only in a single direction:
feature = 2
result, grad, did_finish = eval_diff_tree_array(expression, X, operators, feature)
Generic types
Does this work for only scalar operators on real numbers, or will it work for
MyCrazyType?
Iâm so glad you asked. DynamicExpressions.jl actually will work for arbitrary types! However, to work on operators other than real scalars, you need to use the GenericOperatorEnum <: AbstractOperatorEnum instead of the normal OperatorEnum. Letâs try it with strings!
x1 = Node(String; feature=1)
This node, will be used to index input data (whatever it may be) with either data[feature] (1D abstract arrays) or selectdim(data, 1, feature) (ND abstract arrays). Letâs now define some operators to use:
my_string_func(x::String) = "Hello $x"
operators = GenericOperatorEnum(;
binary_operators=[*],
unary_operators=[my_string_func],
extend_user_operators=true)
Now, letâs create an expression:
tree = x1 * " World!"
tree(["Hello", "Me?"])
# Hello World!
So indeed it works for arbitrary types. It is a bit slower due to the potential for type instability, but itâs not too bad:
@btime tree(["Hello", "Me?"])
# 1738 ns
Tensors
Does this work for tensors, or even unions of scalars and tensors?
Also yes! Letâs see:
using DynamicExpressions
T = Union{Float64,Vector{Float64}}
c1 = Node(T; val=0.0) # Scalar constant
c2 = Node(T; val=[1.0, 2.0, 3.0]) # Vector constant
x1 = Node(T; feature=1)
# Some operators on tensors (multiple dispatch can be used for different behavior!)
vec_add(x, y) = x .+ y
vec_square(x) = x .* x
# Set up an operator enum:
operators = GenericOperatorEnum(;binary_operators=[vec_add], unary_operators=[vec_square], extend_user_operators=true)
# Construct the expression:
tree = vec_add(vec_add(vec_square(x1), c2), c1)
X = [[-1.0, 5.2, 0.1], [0.0, 0.0, 0.0]]
# Evaluate!
tree(X) # [2.0, 29.04, 3.01]
Note that if an operator is not defined for the particular input, nothing will be returned instead.
This is all still pretty fast, too:
@btime tree(X)
# 2,949 ns
@btime eval(:(vec_add(vec_add(vec_square(X[1]), [1.0, 2.0, 3.0]), 0.0)))
# 115,000 ns
Additional notes
Where would I use this instead of GeneralizedGenerated.jl or RuntimeGeneratedFunctions.jl?
These packages have slightly different use cases. DynamicExpressions.jl is more for exploring a very large space of possible functional forms, without ever caching or compiling individual expressions â it only ever compiles the operators used, and an general evaluation function to traverse (very lightweight) expression trees. Thus, the first time evaluating a particular expression will be just as fast as the second time (assuming youâve already compiled the required operators and traversal code).
DynamicExpressions.jl is less flexible: all expressions need to be stored as binary trees, and potential operators have to be known beforehand. GeneralizedGenerated.jl and RuntimeGeneratedFunctions.jl are for cases where you want to generate and compile entire Julia functions at runtime, and re-use them multiple times. Note that they should be compatible with DynamicExpressions.jl, I think, in cases where you want to generate operators, and then use those operators in generated expressions.
Very interested in hearing peopleâs thoughts, feedback, and ideas. This library would not be here without the help of many contributorsâincluding code, advice, and tips on this very forumâ Kaze Wong, @mkitti, @ngam, @johanbluecreek, @marius311, @Julius_Martensen, @ChrisRackauckas, Charles Fox, @patrick-kidger, @jling, vnikoofard, @Skoffer, @Henrique_Becker, shashi, YingboMa, and probably many others Iâm forgetting. Thank you!
Context: DynamicExpressions.jl is actually the result of SymbolicRegression.jl being split in two: the former is exclusively focused on the underlying expression structure for fast evaluation, and the latter will now only worry about evolution-based symbolic optimization (though this splitting is still in-process). There are multiple reasons for the split:
- Makes it easier for others to build symbolic regression libraries, as they can now import just the data structure and evaluation scheme, rather than the heavyweight SymbolicRegression.jl. I hope that this also will introduce new contributors to improve DynamicExpressions.jl so that we can all benefit.
- Lowers code complexity and increases modularity. It should be easier to explore new interfaces with other Julia libraries.
- Will make it easier to include other types of operators, such as tensors (right now, only scalar operators are available), as the modification doesnât need to immediately work with the entirety of SymbolicRegression.jl.