From 27d7386d7c10505e5fe0d481c396b604aef812f8 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 28 Nov 2025 02:13:45 +0000
Subject: [PATCH 1/2] Initial plan


From 2aa2a311a7b79cc4e92646cbb95b5abd90823698 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Fri, 28 Nov 2025 02:16:20 +0000
Subject: [PATCH 2/2] Add docstring to Repeated class clarifying independent
 parameters

Co-authored-by: findmyway <5612003+findmyway@users.noreply.github.com>
---
 src/julax/layers.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/julax/layers.py b/src/julax/layers.py
index c1c6d58..0cdb0fb 100644
--- a/src/julax/layers.py
+++ b/src/julax/layers.py
@@ -38,6 +38,12 @@ def forward(self, x: PyTree, p: Param, s: State) -> tuple[PyTree, State]:
 
 
 class Repeated(LayerBase):
+    """Applies the same layer n times sequentially.
+
+    Each repetition has independent, separately initialized parameters.
+    This is typical for transformer blocks where each layer has its own weights.
+    """
+
     n: int
     layer: LayerLike