Merge pull request #1508 from SebastianM-C/outputsize

make `outputsize` work with symbolically wrapped lux models

Author: ChrisRackauckas

Diff for: ext/SymbolicsLuxExt.jl

@@ -3,6 +3,7 @@ module SymbolicsLuxExt
 using Lux
 using Symbolics
 using Lux.LuxCore
+using Lux.Random: AbstractRNG, default_rng
 using Symbolics.SymbolicUtils
 
 @static if isdefined(Lux.NilSizePropagation, :recursively_nillify)
@@ -11,9 +12,13 @@ using Symbolics.SymbolicUtils
     end
 end
 
+function LuxCore.outputsize(model::SymbolicUtils.BasicSymbolic{<:LuxCore.AbstractLuxLayer}, x::Symbolics.Arr, rng::AbstractRNG)
+    LuxCore.outputsize(Symbolics.getdefaultval(model), x, rng)
+end
+
 @register_array_symbolic LuxCore.stateless_apply(
     model::LuxCore.AbstractLuxLayer, x::AbstractArray, ps::Union{NamedTuple, <:AbstractVector}) begin
-    size = LuxCore.outputsize(model, Symbolics.wrap(x), LuxCore.Random.default_rng())
+    size = LuxCore.outputsize(model, Symbolics.wrap(x), default_rng())
     eltype = Real
 end
 

Diff for: test/extensions/lux.jl

@@ -11,6 +11,7 @@ using ComponentArrays
     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]
+    Symbolics.@variables sym_model::typeof(model) = model
 
     out_ref = LuxCore.stateless_apply(model, x, ps)
     @test out_ref isa Vector{Float32}
@@ -41,6 +42,12 @@ using ComponentArrays
     @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
+
+    out = LuxCore.stateless_apply(sym_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_model => model, sym_x => x, sym_ca => ca),)))
+    @test out_sub == out_ref
 end
 
 @testset "Chain" begin
@@ -53,6 +60,7 @@ end
     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]
+    Symbolics.@variables sym_model::typeof(model) = model
 
     out_ref = LuxCore.stateless_apply(model, x, ps)
     @test out_ref isa Vector{Float32}
@@ -83,4 +91,10 @@ end
     @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
+
+    out = LuxCore.stateless_apply(sym_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_model => model, sym_x => x, sym_ca => ca),)))
+    @test out_sub == out_ref
 end