flodl 0.7.0

    use super::*;
    use crate::tensor::{
        cuda_device_count, cuda_synchronize, test_device, DType, TensorOptions,
    };

    fn require_multi_gpu() -> bool {
        if !test_device().is_cuda() || cuda_device_count() < 2 {
            return false;
        }
        for i in 0..2 {
            let opts = TensorOptions {
                dtype: DType::Float32,
                device: Device::CUDA(i),
            };
            if Tensor::zeros(&[1], opts).is_err() {
                eprintln!(
                    "Device CUDA({i}) cannot run compute kernels, skipping multi-GPU test"
                );
                return false;
            }
        }
        true
    }

    // -- Cross-device autograd verification ---------------------------------

    #[test]
    fn test_cross_device_autograd_gradient_flow() {
        if !require_multi_gpu() {
            return;
        }

        let opts0 = TensorOptions {
            dtype: DType::Float32,
            device: Device::CUDA(0),
        };
        let opts1 = TensorOptions {
            dtype: DType::Float32,
            device: Device::CUDA(1),
        };

        // Parameters on two different devices
        let w0 = Variable::new(Tensor::ones(&[4, 3], opts0).unwrap(), true);
        let w1 = Variable::new(Tensor::ones(&[4, 3], opts1).unwrap(), true);

        // Input on device 0 (no requires_grad, like training data)
        let input = Variable::new(
            Tensor::ones(&[4, 4], opts0).unwrap(),
            false,
        );

        // Chunk along batch dim: 2 shards of size 2
        let chunks = input.chunk(2, 0).unwrap();
        assert_eq!(chunks.len(), 2);

        // Shard 0: forward on device 0
        let out0 = chunks[0].matmul(&w0).unwrap(); // [2, 3] on dev0

        // Shard 1: move to device 1, forward there, move output back to device 0
        let shard1_dev1 = chunks[1].to_device(Device::CUDA(1)).unwrap();
        let out1_dev1 = shard1_dev1.matmul(&w1).unwrap(); // [2, 3] on dev1
        let out1_dev0 = out1_dev1.to_device(Device::CUDA(0)).unwrap(); // [2, 3] on dev0

        // Gather: cat outputs on device 0
        let gathered = Variable::cat_many(&[&out0, &out1_dev0], 0).unwrap(); // [4, 3]

        // Compute scalar loss
        let loss = gathered.sum().unwrap();

        // Backward
        loss.backward().unwrap();

        // Verify: both parameters received gradients on their own devices
        let grad0 = w0.grad();
        let grad1 = w1.grad();
        assert!(
            grad0.is_some(),
            "w0 on device 0 should have gradient after backward"
        );
        assert!(
            grad1.is_some(),
            "w1 on device 1 should have gradient after backward"
        );

        // Verify gradients are on the correct devices
        let g0 = grad0.unwrap();
        let g1 = grad1.unwrap();
        assert_eq!(g0.device(), Device::CUDA(0), "w0 gradient should be on device 0");
        assert_eq!(g1.device(), Device::CUDA(1), "w1 gradient should be on device 1");

        // Verify gradient values are non-zero
        let g0_sum = g0.sum().unwrap().item().unwrap();
        let g1_sum = g1.sum().unwrap().item().unwrap();
        assert!(
            g0_sum.abs() > 1e-6,
            "w0 gradient should be non-zero, got {g0_sum}"
        );
        assert!(
            g1_sum.abs() > 1e-6,
            "w1 gradient should be non-zero, got {g1_sum}"
        );

        cuda_synchronize(0);
        cuda_synchronize(1);
    }

    #[test]
    fn test_cross_device_autograd_values() {
        // Verify that cross-device backward produces the SAME gradients
        // as single-device backward (correctness check).
        if !require_multi_gpu() {
            return;
        }

        // Use deterministic values
        let w_data = Tensor::from_f32(
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
            &[4, 2],
            Device::CUDA(0),
        )
        .unwrap();

        // Single-device reference: forward all on device 0
        let w_ref = Variable::new(w_data.clone(), true);
        let x = Tensor::from_f32(
            &[1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
            &[4, 4],
            Device::CUDA(0),
        )
        .unwrap();
        let x_var = Variable::new(x.clone(), false);
        let out_ref = x_var.matmul(&w_ref).unwrap();
        let loss_ref = out_ref.sum().unwrap();
        loss_ref.backward().unwrap();
        let grad_ref = w_ref.grad().unwrap();
        let grad_ref_vals = grad_ref.to_f32_vec().unwrap();

        // Cross-device: split batch across 2 devices.
        // Create w0 and w1 from fresh tensors (not clones of w_data,
        // which was tainted by set_requires_grad through w_ref's shallow clone).
        let w0 = Variable::new(
            Tensor::from_f32(
                &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
                &[4, 2],
                Device::CUDA(0),
            )
            .unwrap(),
            true,
        );
        let w1 = Variable::new(
            Tensor::from_f32(
                &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
                &[4, 2],
                Device::CUDA(1),
            )
            .unwrap(),
            true,
        );
        let x_var2 = Variable::new(x, false);

        let chunks = x_var2.chunk(2, 0).unwrap();

        let out0 = chunks[0].matmul(&w0).unwrap();
        let shard1 = chunks[1].to_device(Device::CUDA(1)).unwrap();
        let out1_dev1 = shard1.matmul(&w1).unwrap();
        let out1_dev0 = out1_dev1.to_device(Device::CUDA(0)).unwrap();
        let gathered = Variable::cat_many(&[&out0, &out1_dev0], 0).unwrap();
        let loss = gathered.sum().unwrap();
        loss.backward().unwrap();

        // Sum of cross-device gradients should equal single-device gradient
        let g0 = w0.grad().unwrap().to_f32_vec().unwrap();
        let g1 = w1.grad().unwrap().to_f32_vec().unwrap();

        for i in 0..g0.len() {
            let cross_sum = g0[i] + g1[i];
            let diff = (cross_sum - grad_ref_vals[i]).abs();
            assert!(
                diff < 1e-5,
                "gradient mismatch at index {i}: cross-device sum {cross_sum} vs reference {}",
                grad_ref_vals[i]
            );
        }

        cuda_synchronize(0);
        cuda_synchronize(1);
    }

    // -- Graph integration tests (CPU, single-GPU fallback) -----------------

    #[test]
    fn test_graph_set_optimizer_and_step() {
        use crate::graph::FlowBuilder;
        use crate::nn::{Adam, Linear, ReLU, mse_loss};

        let model = FlowBuilder::from(Linear::new(4, 8).unwrap())
            .through(ReLU::new())
            .through(Linear::new(8, 2).unwrap())
            .build()
            .unwrap();

        model.set_optimizer(|p| Adam::new(p, 0.01));
        model.set_training(true);

        // Snapshot initial params
        let params_before: Vec<f32> = model
            .parameters()
            .iter()
            .flat_map(|p| p.variable.data().to_f32_vec().unwrap())
            .collect();

        // One training step
        let x = Variable::new(
            Tensor::randn(&[4, 4], Default::default()).unwrap(),
            false,
        );
        let target = Variable::new(
            Tensor::randn(&[4, 2], Default::default()).unwrap(),
            false,
        );
        let out = model.forward(&x).unwrap();
        let loss = mse_loss(&out, &target).unwrap();
        loss.backward().unwrap();
        model.step().unwrap();

        // Params should have changed
        let params_after: Vec<f32> = model
            .parameters()
            .iter()
            .flat_map(|p| p.variable.data().to_f32_vec().unwrap())
            .collect();

        let changed = params_before
            .iter()
            .zip(&params_after)
            .any(|(a, b)| (a - b).abs() > 1e-8);
        assert!(changed, "parameters should change after step()");
    }

    #[test]
    fn test_graph_step_without_optimizer() {
        use crate::graph::FlowBuilder;
        use crate::nn::Linear;

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

        // step() without set_optimizer() should be a no-op, not a crash
        let result = model.step();
        assert!(result.is_ok());
    }

    #[test]
    fn test_graph_set_lr() {
        use crate::graph::FlowBuilder;
        use crate::nn::{Adam, Linear};

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

        model.set_optimizer(|p| Adam::new(p, 0.01));
        // Should not panic
        model.set_lr(0.001);
    }

    // -- El Che unit tests (CPU, no NCCL needed) ----------------------------

    #[test]
    fn test_cadence_initial_equal() {
        let c = ElChe::new(2, 10);
        assert_eq!(c.batches(0), 10);
        assert_eq!(c.batches(1), 10);
        assert_eq!(c.total_batches(), 20);
        assert_eq!(c.anchor(), 10);
        assert!(!c.is_calibrated());
    }

    #[test]
    fn test_cadence_initial_three_devices() {
        let c = ElChe::new(3, 15);
        assert_eq!(c.batches(0), 15);
        assert_eq!(c.batches(1), 15);
        assert_eq!(c.batches(2), 15);
        assert_eq!(c.total_batches(), 45);
    }

    #[test]
    fn test_cadence_ratio_discovery_2x() {
        // Device 0 is 2x faster than device 1.
        // Equal counts (10:10), device 0 finishes in 500ms, device 1 in 1000ms.
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.50); // high target to avoid anchor auto-tune
        let bc = c.batch_counts().to_vec(); c.report_timing(&[500.0, 1000.0], &bc, 10.0);

        assert!(c.is_calibrated());
        // Slow device (rank 1) keeps anchor=10, fast device (rank 0) gets ~20.
        assert_eq!(c.batches(1), 10);
        assert_eq!(c.batches(0), 20);
    }

    #[test]
    fn test_cadence_ratio_discovery_fbrl_like() {
        // Simulates RTX 5060 Ti vs GTX 1060 (~2.3:1 speed ratio).
        // Anchor=10 on slow device, equal initial counts.
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.50); // no auto-tune

        // Both ran 10 batches; fast took 730ms (73ms/batch), slow took 1640ms (164ms/batch).
        let bc = c.batch_counts().to_vec(); c.report_timing(&[730.0, 1640.0], &bc, 50.0);

        assert!(c.is_calibrated());
        assert_eq!(c.batches(1), 10); // slow device: anchor
        // Fast device: 164/73 * 10 ≈ 22.5, rounds to 22 or 23
        let fast = c.batches(0);
        assert!(
            (22..=23).contains(&fast),
            "expected ~22-23, got {fast}"
        );
    }

    #[test]
    fn test_cadence_anchor_auto_tune() {
        // High per-window fixed overhead should trigger anchor growth.
        // No fill is staged (no coordinator), so the signal is reduce-only:
        // overhead = sync / (anchor·marginal + sync). With marginal 100ms,
        // anchor 10 => window_compute 1000, sync 500 => 500/1500 = 0.33 > 0.10.
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.10);

        // Prime with five low-overhead reports of equal-speed timings (below
        // target, so no growth fires during warmup either), then issue the
        // high-overhead trigger from Stable.
        for _ in 0..5 {
            let bc = c.batch_counts().to_vec();
            c.report_timing(&[1000.0, 1000.0], &bc, 5.0);
        }
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[1000.0, 1000.0], &bc, 500.0);

        // window-pressure proposes; commit to apply (Stable verdict).
        c.commit_proposed_anchor();

        // scale = min(0.33/0.10, GROWTH_STEP_CAP=2.0) = 2.0
        // new anchor = ceil(10 * 2) = 20 (capped per cycle; multi-cycle climb)
        assert_eq!(c.anchor(), 20);
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 20);
    }

    #[test]
    fn test_cadence_anchor_auto_tune_with_speed_ratio() {
        // Heterogeneous: fast device 2x, high sync overhead.
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.10);

        // Prime to Stable phase. Pass fixed bc=[10,10] each call so the
        // synthetic wall_ms keeps a stable per-batch ratio across reports
        // (in production wall_ms would scale with n; in the test it does
        // not, so we keep n fixed instead).
        for _ in 0..5 {
            c.report_timing(&[500.0, 1000.0], &[10, 10], 5.0);
        }
        c.report_timing(&[500.0, 1000.0], &[10, 10], 400.0);
        c.commit_proposed_anchor();

        // anchor rank = slow rank 1 (100ms/batch). window_compute = 10·100 = 1000,
        // sync 400 => overhead 400/1400 = 0.286. scale = min(2.86, 2.0) = 2.0.
        // new anchor = ceil(10 * 2) = 20.
        assert_eq!(c.anchor(), 20);
        assert_eq!(c.batches(1), 20); // slow device
        // fast device: 100ms/batch vs 50ms/batch => 2x ratio => 40
        assert_eq!(c.batches(0), 40);
    }

    #[test]
    fn test_cap_binding_suppresses_anchor_growth() {
        // Anchor wind-up guard: when the window cap is binding (counts
        // scaled down to fit the epoch), measured overhead stays above
        // target forever — growth proposals must be suppressed, or the
        // anchor ratchets toward max_anchor while delivered counts stay
        // pinned and every anchor-derived quantity lies.
        let mut c = ElChe::new(2, 10).with_overhead_target(0.10);
        c.set_max_total_batches(30); // binding almost immediately

        for _ in 0..5 {
            c.report_timing(&[500.0, 1000.0], &[10, 10], 5.0);
        }
        // Force the cap to bind once.
        c.report_timing(&[500.0, 1000.0], &[10, 10], 400.0);
        c.commit_proposed_anchor();
        let anchor_after_first = c.anchor();

        // Keep reporting pathological overhead: with the cap binding, no
        // further growth proposals may land.
        for _ in 0..10 {
            c.report_timing(&[500.0, 1000.0], &[10, 10], 400.0);
            c.commit_proposed_anchor();
        }
        assert!(
            c.anchor() <= anchor_after_first,
            "anchor ratcheted under a binding window cap: {} -> {}",
            anchor_after_first,
            c.anchor(),
        );
        let total = c.batches(0) + c.batches(1);
        assert!(total <= 30, "cap still enforced: total={total}");
    }

    #[test]
    fn test_speed_ratio_clamped_against_degenerate_sample() {
        // One legitimate-but-tiny reading (sub-ms wall over many batches)
        // must not blow the schedule up by a 1e4x ratio on paths without
        // max_batch_diff / max_total_batches.
        let mut c = ElChe::new(2, 10);
        for _ in 0..6 {
            // rank 0: absurdly fast reading; rank 1: 100 ms/batch.
            c.report_timing(&[0.1, 1000.0], &[10, 10], 1.0);
        }
        // ratio would be 10_000x unclamped; with the 64x clamp the fast
        // rank gets at most anchor * 64 batches.
        assert!(
            c.batches(0) <= 10 * 64,
            "ratio clamp failed: fast rank got {} batches",
            c.batches(0),
        );
    }

    #[test]
    fn test_warmup_unsticks_when_pinned_anchor_never_reports() {
        // A pinned anchor rank that never produces a valid reading must
        // not freeze the controller in Warmup forever: after a full trust
        // window of misses, election falls back to ranks with data.
        let mut c = ElChe::new(2, 10).with_initial_anchor(1);
        for _ in 0..10 {
            // rank 1 (the pinned anchor) reports nothing valid; rank 0
            // reports steadily.
            c.report_timing(&[100.0, 0.0], &[10, 0], 1.0);
        }
        assert!(
            c.is_calibrated(),
            "controller stayed un-calibrated: pinned dead anchor froze Warmup",
        );
    }

    #[test]
    fn test_nudge_anchor_down_ignores_nan_factor() {
        let mut c = ElChe::new(2, 10);
        for _ in 0..6 {
            c.report_timing(&[100.0, 100.0], &[10, 10], 1.0);
        }
        let before = c.anchor();
        c.nudge_anchor_down(f64::NAN);
        assert_eq!(c.anchor(), before, "NaN factor must be a no-op");
        c.nudge_anchor_down(0.5);
        assert!(c.anchor() < before, "finite factor still nudges");
    }

    #[test]
    fn test_cadence_window_capped_to_max_total() {
        // Window cap (set by the cluster coordinator to the epoch's batch
        // count): the overhead auto-tune may grow the schedule to amortize
        // an expensive sync, but `recompute_batch_counts` must scale the
        // per-rank counts down proportionally so their sum never exceeds
        // the cap — a reduce window must fit within one epoch. Mirrors
        // `test_cadence_anchor_auto_tune_with_speed_ratio` (which grows to
        // [80, 40], sum 120) but with the total capped at 60.
        let mut c = ElChe::new(2, 10).with_overhead_target(0.10);
        c.set_max_total_batches(40);

        for _ in 0..5 {
            c.report_timing(&[500.0, 1000.0], &[10, 10], 5.0);
        }
        c.report_timing(&[500.0, 1000.0], &[10, 10], 400.0);
        c.commit_proposed_anchor();

        // Uncapped (×2 grow) this is [40, 20] (sum 60). Capped to 40 and
        // scaled proportionally: ~[26, 13] (sum <= 40, ~2x ratio preserved).
        let total = c.batches(0) + c.batches(1);
        assert!(total <= 40, "window capped to max_total: total={total} (<= 40)");
        assert!(
            c.batches(0) > c.batches(1),
            "speed ratio preserved after cap: fast={} slow={}",
            c.batches(0),
            c.batches(1),
        );
    }

    #[test]
    fn test_cadence_anchor_capped_at_max() {
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.01)
            .with_max_anchor(15);

        // Prime to Stable phase before triggering auto-tune.
        for _ in 0..5 {
            let bc = c.batch_counts().to_vec();
            c.report_timing(&[100.0, 100.0], &bc, 0.5);
        }
        // Extreme overhead: sync dominates.
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[100.0, 100.0], &bc, 500.0);
        c.commit_proposed_anchor();

        // ×2 grow wants ceil(10*2)=20, but max_anchor clamps it to 15.
        assert_eq!(c.anchor(), 15);
        assert_eq!(c.batches(0), 15);
    }

    #[test]
    fn test_cadence_stable_when_overhead_low() {
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.10);

        // sync=5ms on 1000ms compute => 0.5% overhead, well below 10%.
        let bc = c.batch_counts().to_vec(); c.report_timing(&[1000.0, 1000.0], &bc, 5.0);

        assert_eq!(c.anchor(), 10); // no change
    }

    #[test]
    fn test_overhead_proposal_committed_on_stable_verdict() {
        // High-overhead trigger should propose a grow; commit applies it.
        let mut c = ElChe::new(2, 10).with_overhead_target(0.10);
        for _ in 0..5 {
            c.report_timing(&[1000.0, 1000.0], &[10, 10], 5.0);
        }
        // overhead = 500/1500 = 0.33 > 0.10; scale = min(3.3, 2.0) = 2.0
        c.report_timing(&[1000.0, 1000.0], &[10, 10], 500.0);
        // Before commit, anchor unchanged.
        assert_eq!(c.anchor(), 10, "report_timing must not mutate anchor");
        c.commit_proposed_anchor();
        assert_eq!(c.anchor(), 20, "commit applies the ×2-capped grow");
    }

    #[test]
    fn test_overhead_grow_vetoed_on_suppress_growth() {
        // Grow proposal + SuppressGrowth verdict → anchor stays put.
        let mut c = ElChe::new(2, 10).with_overhead_target(0.10);
        for _ in 0..5 {
            c.report_timing(&[1000.0, 1000.0], &[10, 10], 5.0);
        }
        c.report_timing(&[1000.0, 1000.0], &[10, 10], 500.0);
        c.veto_proposed_growth();
        assert_eq!(c.anchor(), 10, "SuppressGrowth vetoes the grow proposal");
    }

    #[test]
    fn test_growth_latched_off_after_suppress_growth() {
        // SuppressGrowth latches growth OFF; it re-arms only after
        // GROWTH_REARM_STABLE (5) consecutive Stable verdicts — the margin
        // to the convergence cliff. The controller must not re-attempt
        // growth every cycle while latched.
        let mut c = ElChe::new(2, 10).with_overhead_target(0.10);
        for _ in 0..5 {
            c.report_timing(&[1000.0, 1000.0], &[10, 10], 5.0);
        }
        // Grow proposed, then SuppressGrowth: drop it AND latch growth off.
        c.report_timing(&[1000.0, 1000.0], &[10, 10], 500.0);
        c.veto_proposed_growth();
        assert_eq!(c.anchor(), 10, "SuppressGrowth vetoes the grow");
        assert!(!c.growth_enabled(), "growth latched off");

        // While latched, a high-overhead report proposes nothing → no grow
        // even on a Stable commit. Five Stable commits re-arm the latch.
        for _ in 0..5 {
            c.report_timing(&[1000.0, 1000.0], &[10, 10], 500.0);
            c.commit_proposed_anchor();
            assert_eq!(c.anchor(), 10, "no growth while latched off / re-arming");
        }
        assert!(c.growth_enabled(), "5 consecutive Stable verdicts re-arm growth");

        // A fresh high-overhead report now grows again.
        c.report_timing(&[1000.0, 1000.0], &[10, 10], 500.0);
        c.commit_proposed_anchor();
        assert!(c.anchor() > 10, "growth resumes after re-arm");
    }

    #[test]
    fn test_overhead_proposal_discarded_on_nudge_down() {
        // Grow proposal + NudgeDown verdict → proposal dropped, nudge
        // operates on the current (pre-proposal) anchor.
        let mut c = ElChe::new(2, 20).with_overhead_target(0.10);
        for _ in 0..5 {
            c.report_timing(&[1000.0, 1000.0], &[20, 20], 5.0);
        }
        c.report_timing(&[1000.0, 1000.0], &[20, 20], 500.0);
        // Proposal: ×2-capped grow to ceil(20 * 2.0) = 40. NudgeDown discards
        // that and applies factor 0.5 to the current anchor (20).
        c.discard_proposed_anchor();
        c.nudge_anchor_down(0.5);
        assert_eq!(c.anchor(), 10, "nudge halves the pre-proposal anchor");
    }

    #[test]
    fn test_cadence_three_devices_mixed_speed() {
        let mut c = ElChe::new(3, 10)
            .with_overhead_target(0.50); // no auto-tune

        // Device 0: 3x fast (333ms), device 1: 2x fast (500ms), device 2: slow (1000ms).
        let bc = c.batch_counts().to_vec(); c.report_timing(&[333.0, 500.0, 1000.0], &bc, 10.0);

        assert_eq!(c.batches(2), 10); // slow: anchor
        // Device 1: 100ms/batch vs 33.3ms/batch for device 0
        // Device 0: ratio 100/33.3 = 3.0 => 30
        // Device 1: ratio 100/50 = 2.0 => 20
        assert_eq!(c.batches(0), 30);
        assert_eq!(c.batches(1), 20);
    }

    #[test]
    fn test_cadence_successive_reports_refine() {
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.50);

        // First report: 2x speed ratio.
        let bc = c.batch_counts().to_vec(); c.report_timing(&[500.0, 1000.0], &bc, 10.0);
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 10);

        // Second report: new counts, faster device did 20 in 1000ms (50ms/batch),
        // slow did 10 in 1000ms (100ms/batch). Ratio stays 2:1.
        let bc = c.batch_counts().to_vec(); c.report_timing(&[1000.0, 1000.0], &bc, 10.0);
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 10);
    }

    #[test]
    fn test_callback_slack_reduces_firing_rank_count() {
        // Calibrate: rank 0 fast (50 ms/batch), rank 1 slow (100 ms/batch).
        // Ratio 2:1 → rank 0 gets 20 batches, rank 1 gets 10 (anchor).
        // wall_ms must scale with `bc` to keep ms-per-batch stable in
        // the trust window (50 vs 25 with different bc → different
        // smoothed ratio, which is fine in production but obscures
        // slack arithmetic in a unit test).
        let mut c = ElChe::new(2, 10).with_overhead_target(0.50);
        c.report_timing(&[500.0, 1000.0], &[10, 10], 10.0);
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 10);

        // Stage 200 ms of callback slack on rank 0. At 50 ms/batch
        // that is 4 batches' worth.
        c.apply_callback_slack(&[200.0, 0.0]);

        // Recompute with timings that keep ms/batch at 50/100:
        // rank 0 doing 20 batches in 1000ms; rank 1 doing 10 in 1000ms.
        c.report_timing(&[1000.0, 1000.0], &[20, 10], 10.0);
        assert_eq!(c.batches(0), 16);
        assert_eq!(c.batches(1), 10);

        // Slack auto-clears: next recompute returns rank 0 to 20 (the
        // un-slacked target). Dead-zone hysteresis allows the jump from
        // 16 to 20 since the delta is > 5% of 16.
        c.report_timing(&[800.0, 1000.0], &[16, 10], 10.0);
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 10);
    }

    #[test]
    fn test_callback_slack_clamps_at_one() {
        // Pathologically large slack must not starve the rank entirely.
        let mut c = ElChe::new(2, 10).with_overhead_target(0.50);
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[500.0, 1000.0], &bc, 10.0);
        assert_eq!(c.batches(0), 20);

        // 10s of slack on a 50ms/batch rank = 200 batches, far above
        // the rank's 20-batch quota. Target should clamp at 1.
        c.apply_callback_slack(&[10_000.0, 0.0]);
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[500.0, 1000.0], &bc, 10.0);
        assert_eq!(c.batches(0), 1, "slack must clamp at 1, not starve to 0");
        assert_eq!(c.batches(1), 10);
    }

    #[test]
    fn test_callback_slack_size_mismatch_is_noop() {
        // Wrong-length slack vectors are silently ignored — a misconfig
        // should never crash a running cluster, just leave behavior at
        // the unslacked baseline.
        let mut c = ElChe::new(2, 10).with_overhead_target(0.50);
        c.report_timing(&[500.0, 1000.0], &[10, 10], 10.0);
        assert_eq!(c.batches(0), 20);

        // Length 3 on a 2-rank cluster: rejected.
        c.apply_callback_slack(&[200.0, 0.0, 0.0]);
        assert_eq!(c.pending_callback_slack_ms(), &[0.0, 0.0]);

        // Length 0: also rejected.
        c.apply_callback_slack(&[]);
        assert_eq!(c.pending_callback_slack_ms(), &[0.0, 0.0]);

        // Behavior unchanged after the no-op set. Use stable timings.
        c.report_timing(&[1000.0, 1000.0], &[20, 10], 10.0);
        assert_eq!(c.batches(0), 20);
    }

    #[test]
    fn test_callback_slack_multi_rank() {
        // Slack on multiple ranks simultaneously: each rank's reduction
        // is independent. ms/batch: rank 0 = 33.3ms, rank 1 = 50ms,
        // rank 2 = 100ms (anchor). Targets: rank 0 = 10 * (100/33.3) ≈ 30,
        // rank 1 = 10 * (100/50) = 20, rank 2 = 10.
        let mut c = ElChe::new(3, 10).with_overhead_target(0.50);
        c.report_timing(&[333.0, 500.0, 1000.0], &[10, 10, 10], 10.0);
        let baseline_0 = c.batches(0);
        let baseline_1 = c.batches(1);
        let baseline_2 = c.batches(2);

        // 100ms slack on rank 0 (= 4 batches @ ceil(100/33) = 4)
        // 100ms slack on rank 1 (= 2 batches @ 50ms)
        c.apply_callback_slack(&[100.0, 100.0, 0.0]);
        // Keep ms/batch stable: wall_ms = bc * ms_per_batch.
        // rank 0: 30 * 33.3 ≈ 999; rank 1: 20 * 50 = 1000; rank 2: 10 * 100 = 1000.
        c.report_timing(
            &[baseline_0 as f64 * 33.3, baseline_1 as f64 * 50.0, 1000.0],
            &[baseline_0, baseline_1, baseline_2],
            10.0,
        );

        // Rank 0 drops by 3-4 (33ms/batch → ceil(100/33) = 4).
        assert!(
            c.batches(0) == baseline_0 - 4 || c.batches(0) == baseline_0 - 3,
            "rank 0 expected baseline-3 or baseline-4, got {} (baseline {baseline_0})",
            c.batches(0),
        );
        // Rank 1 drops by 2.
        assert_eq!(c.batches(1), baseline_1 - 2);
        // Rank 2 (no slack) unchanged.
        assert_eq!(c.batches(2), baseline_2);
    }

    #[test]
    fn test_cadence_clamp_total() {
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.50);

        // Fast device gets 20, slow gets 10. Total = 30.
        let bc = c.batch_counts().to_vec(); c.report_timing(&[500.0, 1000.0], &bc, 10.0);

        // Only 15 batches remain in the epoch.
        let clamped = c.clamp_total(15);
        assert_eq!(clamped.iter().sum::<usize>(), 15);
        // Proportions roughly preserved (2:1).
        assert!(clamped[0] >= clamped[1], "fast device should still get more");
    }

    #[test]
    fn test_cadence_clamp_total_no_op_when_within() {
        let c = ElChe::new(2, 10);
        // Total is 20, max is 30 => no clamping needed.
        let clamped = c.clamp_total(30);
        assert_eq!(clamped, vec![10, 10]);
    }

    #[test]
    fn test_cadence_builders() {
        let c = ElChe::new(2, 10)
            .with_overhead_target(0.20)
            .with_max_anchor(100);
        assert_eq!(c.anchor(), 10);
        assert!(!c.is_calibrated());

        // Overhead target clamped to valid range
        let c2 = ElChe::new(2, 5)
            .with_overhead_target(0.001); // below min 0.01
        // Would be clamped to 0.01 internally
        let _ = c2;
    }

    #[test]
    fn test_cadence_max_batch_diff() {
        let c = ElChe::new(2, 10).with_max_batch_diff(5);
        assert_eq!(c.max_batch_diff(), Some(5));

        let c2 = ElChe::new(2, 10);
        assert_eq!(c2.max_batch_diff(), None);
    }

    #[test]
    fn test_batch_count_clamped_to_max_diff() {
        // Setup: 2 GPUs, anchor=10, max_batch_diff=3.
        let mut c = ElChe::new(2, 10).with_max_batch_diff(3);

        // First report (calibration): GPU 0 slow (10ms/batch), GPU 1 fast (2ms/batch).
        // batch_counts are [10, 10] initially, so wall = ms_per_batch * count.
        // GPU 0: 10 batches * 10ms = 100ms. GPU 1: 10 batches * 2ms = 20ms.
        let bc = c.batch_counts().to_vec(); c.report_timing(&[100.0, 20.0], &bc, 0.0);
        assert!(c.is_calibrated());
        // Calibration pass: no clamping. GPU 1 gets 50 batches (ratio 10/2 * 10).
        let counts_after_cal = c.batch_counts().to_vec();
        assert_eq!(counts_after_cal[0], 10);
        assert_eq!(counts_after_cal[1], 50);

        // Second report: GPU 1 suddenly slows to near GPU 0 speed.
        // batch_counts now [10, 50]. GPU 0: 10*10ms=100ms. GPU 1: 50*9ms=450ms.
        // ms_per_batch[1] EMA: alpha=clamp(|9-2|/2, 0.1, 0.8)=0.8, new=0.8*9+0.2*2=7.6
        // slow_ms = max(10, 7.6) = 10. target[1] = 10*(10/7.6)=13.
        // Without clamping: 50 -> 13 (drop of 37). With max_batch_diff=3: 50 -> 47.
        let bc = c.batch_counts().to_vec(); c.report_timing(&[100.0, 450.0], &bc, 0.0);
        let counts = c.batch_counts();
        assert!(counts[1] >= counts_after_cal[1] - 3,
            "batch count drop should be clamped to 3, was {} now {}",
            counts_after_cal[1], counts[1]);
    }

    #[test]
    fn test_cadence_weighted_allreduce_validation() {
        // Validates that Ddp::weighted_all_reduce_gradients rejects
        // mismatched batch_counts length (tested indirectly via the
        // assertion in ElChe that world_size >= 2).
        let c = ElChe::new(2, 10);
        assert_eq!(c.batch_counts().len(), 2);
    }

    #[test]
    #[should_panic(expected = "El Che requires at least 2 devices")]
    fn test_cadence_requires_two_devices() {
        ElChe::new(1, 10);
    }

    #[test]
    #[should_panic(expected = "anchor must be >= 1")]
    fn test_cadence_requires_positive_anchor() {
        ElChe::new(2, 0);
    }

    #[test]
    fn test_cadence_speed_ratio_2x() {
        // Rank 1 is slow, rank 0 is 2x faster
        let c = ElChe::new(2, 10).with_speed_ratio(1, 2.0);
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 10);
    }

    #[test]
    fn test_cadence_speed_ratio_fbrl() {
        // RTX 5060 Ti (rank 0) ~2.3x faster than GTX 1060 (rank 1)
        let c = ElChe::new(2, 10).with_speed_ratio(1, 2.3);
        assert_eq!(c.batches(0), 23);
        assert_eq!(c.batches(1), 10);
    }

    #[test]
    fn test_cadence_speed_ratio_slow_rank_0() {
        // Rank 0 is the slow one (unusual but valid)
        let c = ElChe::new(2, 10).with_speed_ratio(0, 3.0);
        assert_eq!(c.batches(0), 10);
        assert_eq!(c.batches(1), 30);
    }

    #[test]
    fn test_cadence_speed_ratio_equal() {
        let c = ElChe::new(2, 10).with_speed_ratio(1, 1.0);
        assert_eq!(c.batches(0), 10);
        assert_eq!(c.batches(1), 10);
    }

    #[test]
    fn test_cadence_speed_ratio_three_devices() {
        // Rank 2 is slow, others are 3x faster
        let c = ElChe::new(3, 10).with_speed_ratio(2, 3.0);
        assert_eq!(c.batches(0), 30);
        assert_eq!(c.batches(1), 30);
        assert_eq!(c.batches(2), 10);
    }

    #[test]
    fn test_cadence_speed_ratio_three_devices_mid_slow() {
        // Rank 1 is slow, 0 and 2 are fast
        let c = ElChe::new(3, 10).with_speed_ratio(1, 2.0);
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 10);
        assert_eq!(c.batches(2), 20);
    }

    #[test]
    fn test_cadence_max_anchor_one() {
        // max_anchor=1: minimal cadence, sync after every slow-device batch
        let mut c = ElChe::new(2, 1)
            .with_max_anchor(1)
            .with_speed_ratio(1, 2.0);

        assert_eq!(c.batches(0), 2);
        assert_eq!(c.batches(1), 1);

        // High overhead won't increase anchor past 1
        let bc = c.batch_counts().to_vec(); c.report_timing(&[100.0, 200.0], &bc, 500.0);
        assert_eq!(c.anchor(), 1);
    }

    #[test]
    fn test_nudge_anchor_down() {
        // Need calibrated ElChe so recompute_batch_counts works.
        let mut c = ElChe::new(2, 20)
            .with_overhead_target(0.50); // high target to avoid auto-tune interference
        // Calibrate with 2:1 speed ratio (rank 1 slow).
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[50.0, 100.0], &bc, 0.0);
        assert!(c.is_calibrated());
        assert_eq!(c.anchor(), 20);
        assert_eq!(c.batches(0), 40); // fast rank
        assert_eq!(c.batches(1), 20); // slow rank (anchor)

        // Halve the anchor
        c.nudge_anchor_down(0.5);
        assert_eq!(c.anchor(), 10);
        // Batch counts recomputed proportionally
        assert_eq!(c.batches(0), 20);
        assert_eq!(c.batches(1), 10);
    }

    #[test]
    fn test_nudge_anchor_down_clamped_to_one() {
        // Nudging can go below min_anchor but never below 1.
        let mut c = ElChe::new(2, 5);
        assert_eq!(c.anchor(), 5);

        // factor=0.1 -> ceil(5 * 0.1) = 1
        c.nudge_anchor_down(0.1);
        assert_eq!(c.anchor(), 1, "should clamp to 1");
    }

    #[test]
    fn test_nudge_anchor_down_never_increases() {
        let mut c = ElChe::new(2, 10);
        // factor > 1.0 is clamped to 1.0
        c.nudge_anchor_down(2.0);
        assert_eq!(c.anchor(), 10, "should never increase");
    }

    #[test]
    fn test_cadence_speed_ratio_self_corrects() {
        // Start with wrong guess: say rank 0 is slow, but it's actually fast
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.50)
            .with_speed_ratio(0, 2.0);

        // Wrong: rank 0 gets 10, rank 1 gets 20
        assert_eq!(c.batches(0), 10);
        assert_eq!(c.batches(1), 20);

        // Election can change the anchor only once the balancer enters
        // `Phase::Stable` (≥5 calibrations) — by design, no single noisy
        // reading can flip the initial pick. Feed corrective timings (rank
        // 0 is actually 2x faster) for six reports so the 6th sees Stable
        // on entry and re-elects on the trust window. bc passed verbatim
        // each call so the synthetic per-batch arithmetic stays stable.
        for _ in 0..6 {
            c.report_timing(&[500.0, 2000.0], &[10, 20], 10.0);
        }

        // Self-corrected: rank 1 is slow (anchor), rank 0 gets more
        assert_eq!(c.batches(1), c.anchor());
        assert!(c.batches(0) > c.batches(1), "fast device should get more batches");
    }

    // -- PR 1: Phase machine + tie-band anchor election -----------------------

    use crate::distributed::Phase;

    #[test]
    fn test_phase_starts_at_probe() {
        let c = ElChe::new(3, 10);
        assert_eq!(c.phase(), Phase::Probe);
        assert_eq!(c.anchor_rank(), None);
    }

    #[test]
    fn test_phase_advances_on_first_calibration() {
        let mut c = ElChe::new(3, 10).with_overhead_target(0.50);
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[100.0, 380.0, 395.0], &bc, 10.0);
        assert_eq!(c.phase(), Phase::Warmup);
        assert!(c.anchor_rank().is_some());
    }

    #[test]
    fn test_phase_warmup_to_stable_at_5() {
        let mut c = ElChe::new(2, 10).with_overhead_target(0.50);
        for _ in 0..5 {
            let bc = c.batch_counts().to_vec();
            c.report_timing(&[500.0, 1000.0], &bc, 10.0);
        }
        assert_eq!(c.phase(), Phase::Stable);
    }

    #[test]
    fn test_phase_stable_to_mature_at_20() {
        let mut c = ElChe::new(2, 10).with_overhead_target(0.50);
        for _ in 0..20 {
            let bc = c.batch_counts().to_vec();
            c.report_timing(&[500.0, 1000.0], &bc, 10.0);
        }
        assert_eq!(c.phase(), Phase::Mature);
    }

    #[test]
    fn test_anchor_stable_under_tied_slow_ranks() {
        // The 3-GPU bug case: rank 0 fast (100ms), ranks 1 and 2 within 5%
        // of each other (380 vs 395). Old argmax flapped between 1 and 2 each
        // cycle; tie-band + sticky should pin one and keep it.
        let mut c = ElChe::new(3, 10).with_overhead_target(0.50);

        let bc = c.batch_counts().to_vec();
        c.report_timing(&[100.0, 380.0, 395.0], &bc, 10.0);
        let first = c.anchor_rank().expect("anchor elected");

        // Subsequent cycles with the slowest swapping inside the tie band.
        for (a, b) in &[(390.0, 380.0), (385.0, 388.0), (392.0, 386.0)] {
            let bc = c.batch_counts().to_vec();
            c.report_timing(&[100.0, *a, *b], &bc, 10.0);
            assert_eq!(
                c.anchor_rank(), Some(first),
                "anchor must stay sticky across tied slow-rank fluctuations",
            );
        }
    }

    #[test]
    fn test_anchor_switches_when_clear_winner_emerges() {
        // Outside the cohort band (>15% margin), anchor must follow the real slow.
        let mut c = ElChe::new(3, 10).with_overhead_target(0.50);

        c.report_timing(&[100.0, 400.0, 200.0], &[10, 10, 10], 10.0);
        assert_eq!(c.anchor_rank(), Some(1));

        // Rank 2 becomes clearly slower. Anchor swaps are gated to Stable
        // (≥5 calibrations) so only the 6th call onward sees the new
        // election. Push five corrective reports for the trust window to
        // dominate, then assert. bc fixed to keep ms_per_batch stable.
        for _ in 0..5 {
            c.report_timing(&[100.0, 200.0, 600.0], &[10, 10, 10], 10.0);
        }
        assert_eq!(c.anchor_rank(), Some(2), "real slowdown must be tracked");
    }

    #[test]
    fn test_relax_anchor_up_grows_anchor() {
        let mut c = ElChe::new(2, 10).with_overhead_target(0.50);
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[500.0, 1000.0], &bc, 5.0);
        let before = c.anchor();
        c.relax_anchor_up();
        assert_eq!(c.anchor(), before + 1, "anchor should grow by 1 on relax");
    }

    #[test]
    fn test_relax_anchor_up_capped_by_max_batch_diff() {
        // Ratio 1:3 means at anchor=N, batch_counts=[N, 3N], diff=2N.
        // With max_batch_diff=20, anchor caps at 10 (yielding [10, 30]).
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.50)
            .with_max_batch_diff(20);
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[300.0, 900.0], &bc, 5.0); // 1:3 speed ratio
        // Already at anchor=10, [10,30], diff=20. Next relax would project
        // anchor=11 → [11, 33], diff=22 > 20 → refuse.
        let before = c.anchor();
        c.relax_anchor_up();
        assert_eq!(c.anchor(), before, "relax must refuse when projected diff exceeds cap");
    }

    #[test]
    fn test_relax_anchor_up_capped_by_max_anchor() {
        let mut c = ElChe::new(2, 10)
            .with_overhead_target(0.50)
            .with_max_anchor(11);
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[500.0, 1000.0], &bc, 5.0);
        c.relax_anchor_up();
        assert_eq!(c.anchor(), 11);
        c.relax_anchor_up();
        assert_eq!(c.anchor(), 11, "relax must respect max_anchor");
    }

    #[test]
    fn test_anchor_election_lowest_rank_tiebreak() {
        // No prior anchor (Probe phase first call): with all ranks tied, the
        // deterministic tiebreak picks the lowest-indexed candidate.
        let mut c = ElChe::new(3, 10).with_overhead_target(0.50);
        let bc = c.batch_counts().to_vec();
        c.report_timing(&[100.0, 100.0, 100.0], &bc, 10.0);
        assert_eq!(c.anchor_rank(), Some(0));
    }