Skip to content

Commit

Permalink
Merge tag 'v1.8.0' into cuda-5
Browse files Browse the repository at this point in the history
[Diff since v1.7.0](v1.7.0...v1.8.0)

**Merged pull requests:**
- Create preallocation utility functions for expressions (#114) (@MilesCranmer)
  • Loading branch information
MilesCranmer committed Dec 15, 2024
2 parents 1422e36 + dde9291 commit 6c01ded
Show file tree
Hide file tree
Showing 60 changed files with 2,891 additions and 769 deletions.
11 changes: 6 additions & 5 deletions .github/workflows/CI.yml
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,7 @@ jobs:
- macOS-latest
include:
- os: ubuntu-latest
julia-version: '1.7'
- os: ubuntu-latest
julia-version: '1.6'
julia-version: '1.10'

steps:
- uses: actions/checkout@v2
Expand Down Expand Up @@ -64,7 +62,7 @@ jobs:
additional_tests:
name: test ${{ matrix.test_name }} - ${{ matrix.os }}
runs-on: ${{ matrix.os }}
timeout-minutes: 60
timeout-minutes: 120
strategy:
fail-fast: false
matrix:
Expand All @@ -73,7 +71,7 @@ jobs:
julia-version:
- "1"
test_name:
- "enzyme"
# - "enzyme" # flaky; seems to infinitely compile and fail the CI
- "jet"
steps:
- uses: actions/checkout@v2
Expand All @@ -83,13 +81,16 @@ jobs:
- uses: julia-actions/cache@v1
- uses: julia-actions/julia-buildpkg@v1
- name: Run tests
id: run-tests
continue-on-error: ${{ matrix.test_name == 'enzyme' }}
run: |
julia --color=yes -e 'import Pkg; Pkg.add("Coverage")'
SR_TEST=${{ matrix.test_name }} julia --color=yes --threads=auto --check-bounds=yes --depwarn=yes --code-coverage=user -e 'import Coverage; import Pkg; Pkg.activate("."); Pkg.test(coverage=true)'
julia --color=yes coverage.jl
shell: bash
- name: Coveralls
uses: coverallsapp/github-action@v2
if: steps.run-tests.outcome == 'success'
with:
parallel: true
path-to-lcov: lcov.info
Expand Down
3 changes: 1 addition & 2 deletions .github/workflows/Documentation.yml
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,7 @@ jobs:
- name: Build and deploy
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key
DOCUMENTER_KEY_CAM: ${{ secrets.DAMTP_DEPLOY_KEY }}
run: |
cd docs
julia --project=. make.jl
2 changes: 1 addition & 1 deletion .github/workflows/benchmark_pr.yml
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ jobs:
- uses: actions/checkout@v2
- uses: julia-actions/setup-julia@v1
with:
version: "1.9"
version: "1"
- uses: julia-actions/cache@v1
- name: Extract Package Name from Project.toml
id: extract-package-name
Expand Down
10 changes: 3 additions & 7 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,15 +1,13 @@
name = "DynamicExpressions"
uuid = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
authors = ["MilesCranmer <[email protected]>"]
version = "0.18.6"
version = "1.8.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"
PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
Expand All @@ -35,18 +33,16 @@ DynamicExpressionsZygoteExt = "Zygote"
Bumper = "0.6"
CUDA = "4, 5"
ChainRulesCore = "1"
Compat = "3.37, 4"
DispatchDoctor = "0.4"
Interfaces = "0.3"
LoopVectorization = "0.12"
MacroTools = "0.4, 0.5"
Optim = "0.19, 1"
PackageExtensionCompat = "1"
PrecompileTools = "1"
Reexport = "1"
SymbolicUtils = "0.19, ^1.0.5, 2"
SymbolicUtils = "0.19, ^1.0.5, 2, 3"
Zygote = "0.6"
julia = "1.6"
julia = "1.10"

[extras]
Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
Expand Down
108 changes: 46 additions & 62 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,7 @@ A dynamic expression is a snippet of code that can change throughout runtime - c
3. It then generates specialized [evaluation kernels](https://github.com/SymbolicML/DynamicExpressions.jl/blob/fe8e6dfa160d12485fb77c226d22776dd6ed697a/src/EvaluateEquation.jl#L29-L66) for the space of potential operators.
4. It also generates kernels for the [first-order derivatives](https://github.com/SymbolicML/DynamicExpressions.jl/blob/fe8e6dfa160d12485fb77c226d22776dd6ed697a/src/EvaluateEquationDerivative.jl#L139-L175), using [Zygote.jl](https://github.com/FluxML/Zygote.jl).
5. DynamicExpressions.jl 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](https://github.com/JuliaSymbolics/SymbolicUtils.jl), so you can move your runtime expression into a CAS!
6. It also has import and export functionality with [SymbolicUtils.jl](https://github.com/JuliaSymbolics/SymbolicUtils.jl).


## Example
Expand All @@ -29,18 +28,17 @@ It also has import and export functionality with [SymbolicUtils.jl](https://gith
using DynamicExpressions

operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos])
variable_names = ["x1", "x2"]

x1 = Node{Float64}(feature=1)
x2 = Node{Float64}(feature=2)
x1 = Expression(Node{Float64}(feature=1); operators, variable_names)
x2 = Expression(Node{Float64}(feature=2); operators, variable_names)

expression = x1 * cos(x2 - 3.2)

X = randn(Float64, 2, 100);
expression(X, operators) # 100-element Vector{Float64}
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?
Expand All @@ -53,27 +51,25 @@ First, what happens if we naively use Julia symbols to define and then evaluate
This is quite slow, meaning it will be hard to quickly search over the space of expressions. Let's see how DynamicExpressions.jl compares:

```julia
@btime expression(X, operators)
# 693 ns
@btime expression(X)
# 607 ns
```

Much faster! And we didn't even need to compile it. (Internally, this is calling `eval_tree_array(expression, X, operators)`).
Much faster! And we didn't even need to compile it. (Internally, this is calling `eval_tree_array(expression, X)`).

If we change `expression` dynamically with a random number generator, it will have the same performance:

```julia
@btime begin
expression.op = rand(1:3) # random operator in [+, -, *]
expression(X, operators)
end
# 842 ns
@btime ex(X) setup=(ex = copy(expression); ex.tree.op = rand(1:3) #= random operator in [+, -, *] =#)
# 640 ns
```

Now, let's see the performance if we had hard-coded these expressions:

```julia
f(X) = X[1, :] .* cos.(X[2, :] .- 3.2)
@btime f(X)
# 708 ns
# 629 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:
Expand Down Expand Up @@ -102,49 +98,37 @@ We can also compute gradients with the same speed:
```julia
using Zygote # trigger extension

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

expression = x1 * cos(x2 - 3.2)
```

We can take the gradient with respect to inputs with simply the `'` character:

```julia
grad = expression'(X, operators)
grad = expression'(X)
```

This is quite fast:

```julia
@btime expression'(X, operators)
# 2894 ns
@btime expression'(X)
# 2333 ns
```

and again, we can change this expression at runtime, without loss in performance!

```julia
@btime begin
expression.op = rand(1:3)
expression'(X, operators)
end
# 3198 ns
@btime ex'(X) setup=(ex = copy(expression); ex.tree.op = rand(1:3))
# 2333 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:

```julia
result, grad, did_finish = eval_grad_tree_array(expression, X, operators; variable=false)
```

or with respect to variables, and only in a single direction:

```julia
feature = 2
result, grad, did_finish = eval_diff_tree_array(expression, X, operators, feature)
result, grad, did_finish = eval_grad_tree_array(expression, X; variable=false)
```

## Generic types
Expand All @@ -154,42 +138,37 @@ result, grad, did_finish = eval_diff_tree_array(expression, X, operators, featur
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!

```julia
x1 = Node(String; feature=1)
_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:

```julia
my_string_func(x::String) = "ello $x"
using DynamicExpressions: @declare_expression_operator

operators = GenericOperatorEnum(;
binary_operators=[*],
unary_operators=[my_string_func]
)
```
my_string_func(x::String) = "ello $x"
@declare_expression_operator(my_string_func, 1)

Now, let's extend our operators to work with the
expression types used by `DynamicExpressions.jl`:
operators = GenericOperatorEnum(; binary_operators=[*], unary_operators=[my_string_func])

```julia
@extend_operators operators
x1 = Expression(_x1; operators, variable_names)
```

Now, let's create an expression:

```julia
tree = "H" * my_string_func(x1)
expression = "H" * my_string_func(x1)
# ^ `(H * my_string_func(x1))`

tree(["World!", "Me?"], operators)
expression(["World!", "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:

```julia
@btime tree(["Hello", "Me?"], operators)
# 1738 ns
@btime expression(["Hello", "Me?"])
# 103.105 ns (4 allocations: 144 bytes)
```

## Tensors
Expand All @@ -200,37 +179,42 @@ Also yes! Let's see:

```julia
using DynamicExpressions
using DynamicExpressions: @declare_expression_operator

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

# Enable these operators for DynamicExpressions.jl:
@declare_expression_operator(vec_add, 2)
@declare_expression_operator(vec_square, 1)

# Set up an operator enum:
operators = GenericOperatorEnum(;binary_operators=[vec_add], unary_operators=[vec_square])
@extend_operators operators

# Construct the expression:
tree = vec_add(vec_add(vec_square(x1), c2), c1)
variable_names = ["x1"]
c1 = Expression(Node{T}(; val=0.0); operators, variable_names) # Scalar constant
c2 = Expression(Node{T}(; val=[1.0, 2.0, 3.0]); operators, variable_names) # Vector constant
x1 = Expression(Node{T}(; feature=1); operators, variable_names)

expression = 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, operators) # [2.0, 29.04, 3.01]
expression(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:

```julia
@btime tree(X, operators)
# 2,949 ns
@btime expression(X)
# 461.086 ns (13 allocations: 448 bytes)
@btime eval(:(vec_add(vec_add(vec_square(X[1]), [1.0, 2.0, 3.0]), 0.0)))
# 115,000 ns
```
9 changes: 8 additions & 1 deletion benchmark/benchmarks.jl
Original file line number Diff line number Diff line change
Expand Up @@ -74,11 +74,18 @@ function benchmark_evaluation()
extra_kws...
)
suite[T]["evaluation$(extra_key)"] = @benchmarkable(
[eval_tree_array(tree, X, $operators; turbo=$turbo, $extra_kws...) for tree in trees],
[eval_tree_array(tree, X, $operators; kws...) for tree in trees],
setup=(
X=randn(MersenneTwister(0), $T, 5, $n);
treesize=20;
ntrees=100;
kws=$(
if @isdefined(EvalOptions)
(; eval_options=EvalOptions(; turbo=turbo, extra_kws...))
else
(; turbo, extra_kws...)
end
);
trees=[gen_random_tree_fixed_size(treesize, $operators, 5, $T) for _ in 1:ntrees]
)
)
Expand Down
3 changes: 3 additions & 0 deletions docs/Project.toml
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
[deps]
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
Loading

0 comments on commit 6c01ded

Please sign in to comment.