Add partial_apply

This calls `apply` and only returns the first argument.
LuxDL · Feb 13, 2024 · fd9a91f · fd9a91f
1 parent 051c3b1
commit fd9a91f
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/LuxCore.jl b/src/LuxCore.jl
@@ -114,6 +114,15 @@ Simply calls `model(x, ps, st)`
 """
 apply(model::AbstractExplicitLayer, x, ps, st) = model(x, ps, st)
 
+"""
+    partial_apply(model, x, ps, st)
+
+Calls `apply` and only returns the first argument.
+"""
+function partial_apply(model::AbstractExplicitLayer, x, ps, st)
+    return first(apply(model, x, ps, st))
+end
+
 """
     display_name(layer::AbstractExplicitLayer)