Graph Tree: Composable Subgraph Training

let encoder = FlowBuilder::from(Linear::new(4, 8)?)
    .through(GELU)
    .through(Linear::new(8, 4)?)
    .label("encoder")
    .build()?;

let model = FlowBuilder::from(encoder)  // child "encoder" registered
    .through(Linear::new(4, 2)?)
    .build()?;

assert_eq!(model.tree_children().len(), 1);
assert!(model.child_graph("encoder").is_some());

let child = model.child_graph("encoder").unwrap();
assert!(child.is_composed());  // true -- someone is using us

// Validate paths at build time
assert_eq!(model.validate_path("encoder")?, PathKind::Subgraph);
assert_eq!(model.validate_path("encoder.hidden")?, PathKind::Tag);

// Freeze the entire encoder
model.freeze("encoder")?;
assert!(model.is_frozen("encoder")?);

// Thaw just the scan phase within the encoder
model.thaw("encoder.scan")?;

// Read phase stays frozen
assert!(model.is_frozen("encoder.read")?);

// Phase 1: train only the routing layer, everything else frozen
model.freeze("encoder")?;
// ... train ...

// Phase 2: thaw encoder.scan, keep encoder.read frozen
model.thaw("encoder.scan")?;
// ... train with lower LR on scan ...

let mut optimizer = Adam::with_groups()
    .group(&model.parameters_at("meta")?, 0.001)
    .group(&model.parameters_at("encoder.scan")?, 0.0001)
    // encoder.read is frozen -- not in any group
    .build();

let named = model.named_parameters_at("encoder")?;
// Names like "hidden/weight", "hidden/bias" -- the encoder's own names

// Step 1: Train the encoder standalone
let encoder = FlowBuilder::from(scan_module)
    .tag("scan")
    .through(read_module)
    .tag("read")
    .label("encoder")
    .build()?;

// ... train encoder ...
encoder.save_checkpoint("encoder_v1.fdl.gz")?;

// Step 2: Build a larger model with a fresh encoder
let fresh_encoder = FlowBuilder::from(scan_module_new)
    .tag("scan")
    .through(read_module_new)
    .tag("read")
    .label("encoder")
    .build()?;

let model = FlowBuilder::from(fresh_encoder)
    .through(classifier)
    .build()?;

// Load pre-trained weights into the encoder subgraph
let report = model.load_subgraph_checkpoint("encoder", "encoder_v1.fdl.gz")?;
eprintln!("Loaded {}/{} params", report.loaded.len(), report.loaded.len() + report.missing.len());

// Freeze the read phase (pre-trained, proven)
model.freeze("encoder.read")?;

// After forward pass
model.forward(&input)?;

// Read a child's tagged output (null/nil semantics)
match model.tagged_at("encoder.hidden")? {
    Some(v) => println!("hidden shape: {:?}", v.shape()),
    None => println!("not computed yet"),
}
// Err = path doesn't exist (wiring bug)

// Record into child's observation buffer
model.record_at("encoder.loss", loss_value)?;
model.record_at("encoder.accuracy", acc)?;

// Single flush on the parent flushes the entire tree
model.flush(&[]);

// Read trend from child
let trend = model.trend_at("encoder.loss")?;
println!("encoder loss trend: {:?}", trend.last());

model.record_at("subscan.ce", ce_value)?;
model.record_at("letter.accuracy", acc)?;
model.record_scalar("total_loss", total);

model.flush(&[]);  // flushes parent + subscan + letter

// Monitor sees: total_loss, subscan.ce, letter.accuracy
monitor.log(epoch, t.elapsed(), &model);

// Every epoch: flush parent only
model.flush_local(&[]);

// Every 10 epochs: flush the slow child
if epoch % 10 == 0 {
    model.child_graph("auxiliary").unwrap().flush_local(&[]);
}

// For monitoring, choose what to show:
// - latest_metrics_local() = only this graph's own metrics
// - latest_metrics()       = this graph + all children (tree-recursive)
monitor.log(epoch, t.elapsed(), &model);  // uses latest_metrics() by default

let encoder = FlowBuilder::from(module)
    .tag("_plumbing")       // auto-internal (underscore prefix)
    .through(next)
    .tag("output")          // visible from parent
    .label("encoder")
    .build()?;

let model = FlowBuilder::from(encoder)
    .through(Linear::new(4, 2)?)
    .build()?;

// This fails: _plumbing is internal
assert!(model.tagged_at("encoder._plumbing").is_err());

// This works
assert!(model.tagged_at("encoder.output").is_ok());

FlowBuilder::from(module)
    .tag("intermediate")
    .internal("intermediate")  // explicitly hide from parent
    .through(next)
    .build()?;

// Put encoder in eval mode (BatchNorm uses running stats)
model.set_training_at("encoder", false)?;

// Rest of model stays in training mode

let model = FlowBuilder::from(encoder)
    .through(classifier)
    .verbose(true)
    .build()?;

use flodl::*;

// 1. Build and train the encoder standalone.
let encoder = FlowBuilder::from(scan_module)
    .tag("scan")
    .through(read_module)
    .tag("read")
    .label("encoder")          // the label is what makes it a subgraph
    .build()?;
// ... train encoder ...
encoder.save_checkpoint("encoder_v1.fdl.gz")?;

// 2. Nest a fresh encoder of the same shape inside a larger model.
let model = FlowBuilder::from(fresh_encoder)   // child "encoder" auto-registers
    .through(classifier)
    .label("classifier_v1")
    .build()?;

// 3. Load the pretrained weights into just that subgraph.
let report = model.load_subgraph_checkpoint("encoder", "encoder_v1.fdl.gz")?;
println!("loaded {} params into encoder", report.loaded.len());

// 4. Freeze the proven read phase; let the scan phase keep adapting.
model.freeze("encoder.read")?;

// 5. One optimizer, three learning rates by path.
let mut optimizer = Adam::with_groups()
    .group(&model.parameters_at("classifier_v1")?, 1e-3)
    .group(&model.parameters_at("encoder.scan")?, 1e-4)
    // encoder.read is frozen - deliberately in no group
    .build();

// 6. Observation crosses the boundary: one flush covers the tree.
model.record_at("encoder.loss", enc_loss)?;
model.flush(&[]);
let trend = model.trend_at("encoder.loss")?;

use flodl::nn::{checkpoint_version, migrate_checkpoint_file};

if checkpoint_version("encoder_v1.fdl")? < 2 {
    let report = migrate_checkpoint_file(
        "encoder_v1.fdl",
        "encoder_v2.fdl",
        &encoder.named_parameters(),
        &encoder.named_buffers(),
    )?;
    println!("{}", report);
    assert!(report.is_complete());
}

// Load into the subgraph as usual
model.load_subgraph_checkpoint("encoder", "encoder_v2.fdl")?;