Add dynamic expressions extension

[avik-pal]

• [x] Add a basic working Dynamic Expressions Extension
• [x] AD Support
  • [x] Chain Rules (Zygote)
  • [x] ReverseDiff
  • [x] Tracker
  • [x] ForwardDiff
• [x] Test with SymbolicRegression.jl
• [x] Tests
• [x] Add other options like bumper and loop vectorization in the layer struct
• [x] Add a tutorial on using SymbolicRegression/DynamicExpressions with Lux
• [x] Documentation

This turned out to be easier than I expected.

cc @MilesCranmer

[avik-pal]

using Lux, Random
using DynamicExpressions
using Zygote, Tracker, ReverseDiff
using ComponentArrays

operators = OperatorEnum(; binary_operators=[+, -, *], unary_operators=[cos])

x1 = Node(; feature=1)
x2 = Node(; feature=2)

expr_1 = x1 * cos(x2 - 3.2)
expr_2 = x2 - x1 * x2 + 3.2 - 1.0 * x1

layer = DynamicExpressionsLayer(operators, expr_1, expr_2)

ps, st = Lux.setup(Random.default_rng(), layer)

x = rand(Float32, 2, 16)

layer(x, ps, st)

Zygote.gradient(Base.Fix1(sum, abs2) ∘ first ∘ layer, x, ps, st)

Tracker.gradient(Base.Fix1(sum, abs2) ∘ first ∘ layer, x, ps, st)

ps2 = ComponentArray(ps)

ReverseDiff.gradient((x, ps) -> sum(abs2, first(layer(x, ps, st))), (x, ps2))

[avik-pal]

https://lux.csail.mit.edu/previews/PR585/tutorials/advanced/2_SymbolicOptimalControl is a preview of a tutorial