-
Notifications
You must be signed in to change notification settings - Fork 157
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add array function registration for Lux #1054
Merged
ChrisRackauckas
merged 10 commits into
JuliaSymbolics:master
from
SebastianM-C:luxcore_ext
Mar 3, 2024
Merged
Changes from all commits
Commits
Show all changes
10 commits
Select commit
Hold shift + click to select a range
2f25742
Add extension for LuxCore
SebastianM-C 9dbd42d
add LuxCore extension tests
SebastianM-C 356e4c9
use `LuxCore.outputsize`
SebastianM-C 639a354
Improve tests for the Lux extenstion
SebastianM-C 26ddc9c
update api
SebastianM-C 2577243
bump version
SebastianM-C a6f88bc
remove `AbstractExplicitContainerLayer` specialization
SebastianM-C 78a2732
bump versions
SebastianM-C 5dfa5f5
Merge branch 'master' into luxcore_ext
SebastianM-C d3ebacd
use the more general outputsize
SebastianM-C File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,7 +1,7 @@ | ||
name = "Symbolics" | ||
uuid = "0c5d862f-8b57-4792-8d23-62f2024744c7" | ||
authors = ["Shashi Gowda <[email protected]>"] | ||
version = "5.22.1" | ||
version = "5.23.0" | ||
|
||
[deps] | ||
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" | ||
|
@@ -40,11 +40,13 @@ SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b" | |
|
||
[weakdeps] | ||
Groebner = "0b43b601-686d-58a3-8a1c-6623616c7cd4" | ||
LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" | ||
PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" | ||
SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6" | ||
|
||
[extensions] | ||
SymbolicsGroebnerExt = "Groebner" | ||
SymbolicsLuxCoreExt = "LuxCore" | ||
SymbolicsPreallocationToolsExt = "PreallocationTools" | ||
SymbolicsSymPyExt = "SymPy" | ||
|
||
|
@@ -65,6 +67,7 @@ LaTeXStrings = "1.3" | |
LambertW = "0.4.5" | ||
Latexify = "0.16" | ||
LogExpFunctions = "0.3" | ||
LuxCore = "0.1.11" | ||
MacroTools = "0.5" | ||
NaNMath = "1" | ||
PrecompileTools = "1" | ||
|
@@ -83,8 +86,10 @@ julia = "1.10" | |
|
||
[extras] | ||
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" | ||
ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" | ||
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" | ||
Groebner = "0b43b601-686d-58a3-8a1c-6623616c7cd4" | ||
Lux = "b2108857-7c20-44ae-9111-449ecde12c47" | ||
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" | ||
PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d" | ||
PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" | ||
|
@@ -95,4 +100,4 @@ SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6" | |
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" | ||
|
||
[targets] | ||
test = ["Test", "SafeTestsets", "Pkg", "PkgBenchmark", "PreallocationTools", "ForwardDiff", "Groebner", "BenchmarkTools", "ReferenceTests", "SymPy", "Random"] | ||
test = ["Test", "SafeTestsets", "Pkg", "PkgBenchmark", "PreallocationTools", "ForwardDiff", "Groebner", "BenchmarkTools", "ReferenceTests", "SymPy", "Random", "Lux", "ComponentArrays"] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,11 @@ | ||
module SymbolicsLuxCoreExt | ||
|
||
using LuxCore, Symbolics | ||
|
||
@register_array_symbolic LuxCore.stateless_apply( | ||
model::LuxCore.AbstractExplicitLayer, x::AbstractArray, ps::Union{NamedTuple, <:AbstractVector}) begin | ||
size = LuxCore.outputsize(model, x, LuxCore._default_rng()) | ||
eltype = Real | ||
end | ||
|
||
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,86 @@ | ||
using Symbolics, Lux, Random, Test | ||
using ComponentArrays | ||
|
||
@testset "Dense" begin | ||
model = Dense(5, 6) | ||
rng = Random.default_rng() | ||
x = randn(rng, Float32, 5) | ||
ps, _ = LuxCore.setup(rng, model) | ||
|
||
ca = ComponentArray(ps) | ||
Symbolics.@variables sym_ca[1:length(ca)] = ca | ||
Symbolics.@variables sym_ps::typeof(ps) = ps | ||
Symbolics.@variables sym_x[1:5] = Float32[1,2,3,4,5] | ||
|
||
out_ref = LuxCore.stateless_apply(model, x, ps) | ||
@test out_ref isa Vector{Float32} | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, ps) | ||
# if the symbolic function regstration wouldn't work, we'd get the | ||
# vector propagated through the neural network and get a Vector{Num} | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 6 | ||
# test that we can recover the same value as when using concrete numbers | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x),))) | ||
@test out_sub == out_ref | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, sym_ps) | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 6 | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x, sym_ps => ps),))) | ||
@test out_sub == out_ref | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, ca) | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 6 | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x,),))) | ||
@test out_sub == out_ref | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, sym_ca) | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 6 | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x, sym_ca => ca),))) | ||
@test out_sub == out_ref | ||
end | ||
|
||
@testset "Chain" begin | ||
model = Chain(Dense(5, 6), Dense(6, 2), Dense(2, 3)) | ||
rng = Random.default_rng() | ||
x = randn(rng, Float32, 5) | ||
ps, _ = LuxCore.setup(rng, model) | ||
|
||
ca = ComponentArray(ps) | ||
Symbolics.@variables sym_ca[1:length(ca)] = ca | ||
Symbolics.@variables sym_ps::typeof(ps) = ps | ||
Symbolics.@variables sym_x[1:5] = Float32[1, 2, 3, 4, 5] | ||
|
||
out_ref = LuxCore.stateless_apply(model, x, ps) | ||
@test out_ref isa Vector{Float32} | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, ps) | ||
# if the symbolic function regstration wouldn't work, we'd get the | ||
# vector propagated through the neural network and get a Vector{Num} | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 3 | ||
# test that we can recover the same value as when using concrete numbers | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x),))) | ||
@test out_sub == out_ref | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, sym_ps) | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 3 | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x, sym_ps => ps),))) | ||
@test out_sub == out_ref | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, ca) | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 3 | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x,),))) | ||
@test out_sub == out_ref | ||
|
||
out = LuxCore.stateless_apply(model, sym_x, sym_ca) | ||
@test out isa Symbolics.Arr | ||
@test length(out) == 3 | ||
out_sub = Symbolics.value.(Symbolics.substitute.(Symbolics.scalarize(out), (Dict(sym_x => x, sym_ca => ca),))) | ||
@test out_sub == out_ref | ||
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe also add some tests for Layers where we haven't defined
outputsize
to check that our fallback worksThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This one seems a bit tricky, as the
outputsize
fallback receives a symbolic array that's not anAbstractArray
, so we might need to register additional functions, but that would need to make the extension on Lux and not just LuxCore.I tried changing the extension to Lux and registering
apply
forDense
, but I hitapplyparallel
and it tries to iterateSymbolicUtils.BasicSymbolic{AbstractVector{Real}}
, which fails. I'm not sure if we can continue on this path.@ChrisRackauckas Should I leave the extension for
LuxCore
only or change toLux
? There's no immediate benefit for doing this, but it might make it easier if we wan to support more complex networks in the future, not sure how though, as the fact that we don't have anAbstractArray
seems like a deal breaker.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For now, leave as LuxCore
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ok, in that case this should be good to go. Can you trigger / allow CI?