1 files changed, 75 insertions, 3 deletions

diff --git a/scripts/compressed_operator_predictor.py b/scripts/compressed_operator_predictor.py
index ce52c7a..3f3f73b 100644
--- a/scripts/compressed_operator_predictor.py
+++ b/scripts/compressed_operator_predictor.py
@@ -23,6 +23,9 @@ from fa_tangent_kernel_capacity import pseudo_jacobian  # noqa: E402
 
 @dataclass(frozen=True)
 class CompressedRow:
+    hidden_layers: int
+    width: int
+    train_samples: int
     init_seed: int
     feedback_seed: int
     target_steps: int
@@ -54,6 +57,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--input-dim", type=int, default=16)
     parser.add_argument("--output-dim", type=int, default=4)
     parser.add_argument("--width", type=int, default=64)
+    parser.add_argument("--hidden-layers", type=int, default=2)
     parser.add_argument("--train-samples", type=int, default=128)
     parser.add_argument("--test-samples", type=int, default=512)
     parser.add_argument("--target-steps", type=int, default=50)
@@ -114,6 +118,57 @@ def loss_from_residual(residual: np.ndarray, samples: int) -> float:
     return 0.5 * float(residual @ residual) / samples
 
 
+def initialize_weights(
+    config: dcs.RunConfig,
+    width: int,
+    hidden_layers: int,
+    seed: int,
+) -> list[torch.Tensor]:
+    generator = torch.Generator(device=config.device)
+    generator.manual_seed(seed)
+    dims = [config.input_dim, *([width] * hidden_layers), config.output_dim]
+    weights: list[torch.Tensor] = []
+    for layer, (fan_in, fan_out) in enumerate(zip(dims[:-1], dims[1:])):
+        scale = np.sqrt(2.0 / fan_in) if layer < len(dims) - 2 else 1.0 / np.sqrt(fan_in)
+        weights.append(
+            torch.randn(
+                fan_out,
+                fan_in,
+                generator=generator,
+                dtype=torch.float64,
+                device=config.device,
+            )
+            * scale
+        )
+    return weights
+
+
+def init_feedback(
+    config: dcs.RunConfig,
+    width: int,
+    hidden_layers: int,
+    seed: int,
+) -> list[torch.Tensor]:
+    generator = torch.Generator(device=config.device)
+    generator.manual_seed(seed)
+    shapes = [(width, width)] * max(hidden_layers - 1, 0) + [
+        (width, config.output_dim)
+    ]
+    feedback: list[torch.Tensor] = []
+    for rows, cols in shapes:
+        feedback.append(
+            torch.randn(
+                rows,
+                cols,
+                generator=generator,
+                dtype=torch.float64,
+                device=config.device,
+            )
+            * dcs.feedback_scale(rows, config.feedback_scale)
+        )
+    return feedback
+
+
 def kernel_pair(
     weights: list[torch.Tensor],
     x: torch.Tensor,
@@ -211,14 +266,16 @@ def projection_alpha(k_fa_s: np.ndarray, k_fa0: np.ndarray, k_bp0: np.ndarray) -
 
 def run_one(
     config: dcs.RunConfig,
+    hidden_layers: int,
     x: torch.Tensor,
     y: torch.Tensor,
     init_seed: int,
     feedback_seed: int,
     early_step: int,
 ) -> CompressedRow:
-    initial = dcs.initialize_weights(config, config.widths[0], init_seed)
-    feedback = dcs.init_feedback(config, config.widths[0], feedback_seed)
+    width = config.widths[0]
+    initial = initialize_weights(config, width, hidden_layers, init_seed)
+    feedback = init_feedback(config, width, hidden_layers, feedback_seed)
     with torch.no_grad():
         r0 = (dcs.predict(initial, x) - y).reshape(-1).cpu().numpy()
 
@@ -296,6 +353,9 @@ def run_one(
     retangent_gap = retangent_fa - retangent_bp
 
     return CompressedRow(
+        hidden_layers=hidden_layers,
+        width=width,
+        train_samples=config.train_samples,
         init_seed=init_seed,
         feedback_seed=feedback_seed,
         target_steps=config.steps,
@@ -334,6 +394,10 @@ def write_rows(path: Path, rows: list[CompressedRow]) -> None:
 
 def main() -> None:
     args = parse_args()
+    if args.hidden_layers < 1:
+        raise ValueError("--hidden-layers must be positive.")
+    if any(early_step >= args.target_steps for early_step in args.early_steps):
+        raise ValueError("Every --early-steps value must be smaller than --target-steps.")
     if args.torch_threads > 0:
         torch.set_num_threads(args.torch_threads)
     config = make_config(args)
@@ -354,7 +418,15 @@ def main() -> None:
                     flush=True,
                 )
                 rows.append(
-                    run_one(config, x_train, y_train, init_seed, feedback_seed, early_step)
+                    run_one(
+                        config,
+                        args.hidden_layers,
+                        x_train,
+                        y_train,
+                        init_seed,
+                        feedback_seed,
+                        early_step,
+                    )
                 )
 
     outdir = Path(args.outdir)