use LuxCore.outputsize

JuliaSymbolics · Feb 16, 2024 · 356e4c9 · 356e4c9
1 parent 9dbd42d
commit 356e4c9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/Project.toml b/Project.toml
@@ -67,7 +67,7 @@ LaTeXStrings = "1.3"
 LambertW = "0.4.5"
 Latexify = "0.16"
 LogExpFunctions = "0.3"
-LuxCore = "0.1.7"
+LuxCore = "0.1.8"
 MacroTools = "0.5"
 NaNMath = "1"
 PrecompileTools = "1"

diff --git a/ext/SymbolicsLuxCoreExt.jl b/ext/SymbolicsLuxCoreExt.jl
@@ -3,14 +3,14 @@ module SymbolicsLuxCoreExt
 using LuxCore, Symbolics
 
 @register_array_symbolic LuxCore.partial_apply(
-    model::LuxCore.AbstractExplicitContainerLayer, x::AbstractArray, ps::NamedTuple, st::NamedTuple) begin
-    size = ((model[end].out_dims),)
+    model::LuxCore.AbstractExplicitContainerLayer, x::AbstractArray, ps::Union{NamedTuple, <:AbstractVector}, st::NamedTuple) begin
+    size = LuxCore.outputsize(model[end])
     eltype = Real
 end
 
 @register_array_symbolic LuxCore.partial_apply(
-    model::LuxCore.AbstractExplicitLayer, x::AbstractArray, ps::NamedTuple, st::NamedTuple) begin
-    size = ((model.out_dims),)
+    model::LuxCore.AbstractExplicitLayer, x::AbstractArray, ps::Union{NamedTuple, <:AbstractVector}, st::NamedTuple) begin
+    size = LuxCore.outputsize(model)
     eltype = Real
 end