[ANN] DynamicDiff.jl – Fast Symbolic Differentiation for Runtime-Generated Expressions

DynamicDiff.jl

Dev Build Status Coverage

DispatchDoctor Aqua

I’m happy to announce DynamicDiff.jl, a package that provides extremely fast, higher-order, compilation-free symbolic differentiation of runtime-generated expressions. Built on top of DynamicExpressions.jl, DynamicDiff.jl is designed for scenarios where you continuously generate new symbolic forms and need their derivatives on the fly—without incurring compilation overhead each time. Unlike DynamicExpression.jl’s existing gradient utilities, which are only 1st order, DynamicDiff.jl offers full support for higher-order derivatives!

This library integrates seamlessly with SymbolicRegression.jl, enabling new workflows like “integration by evolution,” where you can discover integrals by ensuring the derivative of an evolved expression matches your target data. Rather than explicitly integrating symbolic expressions or differentiating closed-forms by hand, you simply fit the derivative and let the regression discover the antiderivative.

Key Features

  • Instant Differentiation:
    • No code generation or caching. Differentiation is done by manipulating integer-coded operator trees directly. The only compilation needed is once per operator.
  • High Performance:
    • Even for newly generated expressions, taking a derivative can be on the order of tens of nanoseconds—no warm-up required.
  • Supports Arbitrary Operators:
    • Common operators (sin, cos, +, *, /, etc.) have hand-written derivative rules to enable simplification rules during derivative calculations. Unknown operators use ForwardDiff.jl to generate per-operator derivatives.
  • Integration with SymbolicRegression.jl:
    • Combine D(…) operators with TemplateExpressions in SymbolicRegression.jl to discover integrals or higher-order derivatives as just another step in the symbolic search process.

Example

First, let’s set up some variables with a given set of operators:

using DynamicDiff, DynamicExpressions

operators = OperatorEnum(; binary_operators=(+, *, /, -), unary_operators=(sin, cos));
variable_names = ["x1", "x2", "x3"];
x1, x2, x3 = (Expression(Node{Float64}(feature=i); operators, variable_names) for i in 1:3);

Now, we can generate some symbolic functions and take derivatives:

julia> f = x1 * sin(x2 - 0.5)
x1 * sin(x2 - 0.5)

julia> D(f, 1)
sin(x2 - 0.5)

julia> D(f, 2)
x1 * cos(x2 - 0.5)

These symbolic derivatives are done by simply incrementing integers and arranging a binary tree, so this process is very fast:

julia> using BenchmarkTools

julia> @btime D($f, 2);
  52.865 ns (5 allocations: 240 bytes)

This isn’t compiled or cached! To show this, let’s randomly generate arbitrary expressions and then take derivatives of them:

julia> @btime D(g, 1) setup=(g = [x1, x2, x3][rand(1:3)] * sin([x1, x2, x3][rand(1:3)] - randn())) evals=100
  60.270 ns (2 allocations: 80 bytes)

Example in Symbolic Regression

Here’s how you might use DynamicDiff.jl with SymbolicRegression.jl to discover an integral from an integrand \frac{1}{x^2 \sqrt{x^2 - 1}} in the range x > 1.

using SymbolicRegression, DynamicDiff, MLJ
using DynamicDiff: D

x = 9rand(1000) .+ 1

y = @. 1/(x^2 * sqrt(x^2 - 1))

# Template structure: D(f,1)(x) means we differentiate f(x) wrt x and match that to y.

structure = TemplateStructure{(:f,)}(
    ((; f), (x,)) -> D(f, 1)(x)
)

model = SRRegressor(
    binary_operators=(+, -, *, /),
    unary_operators=(sqrt,),
    maxsize=20,
    expression_type=TemplateExpression,
    expression_options=(; structure),
)

X = (; x)
mach = machine(model, X, y)
fit!(mach)

When the regression completes, it discovers f(x) whose derivative is y. If done correctly, f(x) should find something like \frac{\sqrt{x^2 - 1}}{x}.

How Does It Work?

DynamicDiff.jl treats differentiation as a simple tree transformation. Rather than generating new code or leveraging a symbolic macro system, it:

  1. Expands the Operator Enum:

Each operator (e.g. sin, cos, +, *) corresponds to an integer (via the DynamicExpressions.OperatorEnum). For derivative operators, DynamicDiff.jl extends this operator set at runtime, adding derivative versions. For unary operators, it doubles the set (the original operator plus its derivative). For binary operators, it triples the set (adding partial derivatives w.r.t. each argument).

  1. Integer-Based Tree Manipulation:

Expressions are stored as integer-coded trees referencing these operators. Differentiation means creating a new tree that uses these extended operators, applying known simplifications for trivial cases (like \frac{d}{dx}(f(y)) = 0).

This usually means we only need to perform the transformation

node.op += num_ops

to generate the symbolic derivative (for unary). Binary operators are just node.op += diff_arg == 1 ? num_ops : num_ops * 2

  1. ForwardDiff for Unknowns:

If an operator doesn’t have a known analytic derivative rule, DynamicDiff.jl uses ForwardDiff for the general cases:

function (d::OperatorDerivative{F,1,1})(x) where {F}
    return ForwardDiff.derivative(d.op, x)
end
function (d::OperatorDerivative{F,2,1})(x, y) where {F}
    return ForwardDiff.derivative(Fix{2}(d.op, y), x)
end
function (d::OperatorDerivative{F,2,2})(x, y) where {F}
    return ForwardDiff.derivative(Fix{1}(d.op, x), y)
end

Since the underlying operators we differentiate tend to be pretty simple, this means we can actually get some nice speedups relative to pure ForwardDiff because it otherwise would need to differentiate the entire expression top-to-bottom which can prevent the compiler from performing certain simplifications:

julia> using DynamicExpressions, ForwardDiff, DynamicDiff

julia> operators = OperatorEnum(; binary_operators=[+, -, *, /], unary_operators=[sin, cos]);

julia> variable_names = ["x1", "x2", "x3"];

julia> x1, x2, x3 = (i -> Expression(Node{Float64}(; feature=i); operators, variable_names)).(1:3);

julia> ex = x1 * sin(x2 - 0.5) - x3 * sin(x1) * sin(x2)
(x1 * sin(x2 - 0.5)) - ((x3 * sin(x1)) * sin(x2))

julia> using BenchmarkTools

julia> @btime let fnc=$(y -> ForwardDiff.derivative(let ex=ex; x -> (ex)([x 1.0 2.0]'); end, y))
           ForwardDiff.derivative(fnc, 0.5)  # 2nd order diff
       end
  1.133 μs (36 allocations: 2.19 KiB)
1-element Vector{Float64}:
 0.8068453602226698

julia> @btime D(D($ex, 1), 1)($([0.5 1.0 2.0]'))
  730.825 ns (28 allocations: 1.19 KiB)
1-element Vector{Float64}:
 0.8068453602226698

Eager to hear what people think and feel free to share any feedback!

34 Likes