From fd9a91f33e5517f55e8d2ff10605b1c4f7fbc00f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Wed, 14 Feb 2024 01:58:43 +0200
Subject: [PATCH] Add `partial_apply`

This calls `apply` and only returns the first argument.
---
 src/LuxCore.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/LuxCore.jl b/src/LuxCore.jl
index ae5e66c..7a026b3 100644
--- 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)