Add stateless_apply

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxCore"
 uuid = "bb33d45b-7691-41d6-9220-0943567d0623"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.1.9"
+version = "0.1.10"
 
 [deps]
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"

src/LuxCore.jl

@@ -70,6 +70,11 @@ function initialstates(rng::AbstractRNG, l)
     throw(MethodError(initialstates, (rng, l)))
 end
 
+_getemptystate(::AbstractExplicitLayer) = NamedTuple()
+function _getemptystate(l::NamedTuple{fields}) where {fields}
+    return NamedTuple{fields}(map(_getemptystate, values(l)))
+end
+
 """
     parameterlength(layer)
 
@@ -126,6 +131,15 @@ Simply calls `model(x, ps, st)`
 """
 apply(model::AbstractExplicitLayer, x, ps, st) = model(x, ps, st)
 
+"""
+    stateless_apply(model, x, ps)
+
+Calls `apply` and only returns the first argument.
+"""
+function stateless_apply(model::AbstractExplicitLayer, x, ps)
+    return first(apply(model, x, ps, _getemptystate(model)))
+end
+
 """
     display_name(layer::AbstractExplicitLayer)
 
@@ -179,6 +193,11 @@ function statelength(l::AbstractExplicitContainerLayer{layers}) where {layers}
     return sum(statelength, getfield.((l,), layers))
 end
 
+function _getemptystate(l::AbstractExplicitContainerLayer{layers}) where {layers}
+    length(layers) == 1 && return _getemptystate(getfield(l, first(layers)))
+    return NamedTuple{layers}(_getemptystate.(getfield.((l,), layers)))
+end
+
 # Make AbstractExplicit Layers Functor Compatible
 function Functors.functor(::Type{<:AbstractExplicitContainerLayer{layers}},
         x) where {layers}

test/runtests.jl

@@ -47,6 +47,9 @@ end
 
             @test LuxCore.apply(model, x, ps, st) == model(x, ps, st)
 
+            @test LuxCore.stateless_apply(model, x, ps) ==
+                  first(LuxCore.apply(model, x, ps, NamedTuple()))
+
             @test_nowarn println(model)
         end
 
@@ -88,6 +91,9 @@ end
 
         @test LuxCore.apply(model, x, ps, st) == model(x, ps, st)
 
+        @test LuxCore.stateless_apply(model, x, ps) ==
+              first(LuxCore.apply(model, x, ps, st))
+
         @test_nowarn println(model)
 
         model = Chain2(Dense(5, 5), Dense(5, 6))
@@ -103,6 +109,9 @@ end
 
         @test LuxCore.apply(model, x, ps, st) == model(x, ps, st)
 
+        @test LuxCore.stateless_apply(model, x, ps) ==
+              first(LuxCore.apply(model, x, ps, st))
+
         @test_nowarn println(model)
     end