rlevo_evolution/algorithms/neuroevolution/

arch_nas.rs

//! Bounded architecture NAS — evolving *which* fixed-topology Burn `Module`
//! variant wins a task, alongside that variant's weights.
//!
//! Where [`WeightOnly`](super::weight_only::WeightOnly) evolves the weights of
//! **one** declared topology, this module adds the architecture axis without the
//! full topology-evolution machinery of NEAT: the user declares a small, fixed
//! menu of concrete `Module` variants, and each population member carries a
//! categorical **architecture id** plus a per-variant **weight vector**.
//!
//! # Why a custom harness, not [`Strategy`](crate::strategy::Strategy)
//!
//! [`Strategy<B>`](crate::strategy::Strategy)'s `Genome = Tensor<B, 2>` is a
//! homogeneous float contract. An architecture selector is a *categorical
//! integer*; encoding it as a float is fragile under mutation, and a parallel
//! `Int` tensor breaks the `Strategy` signature. So [`ArchNasStrategy`] is a
//! **custom harness** with its own [`NasGenome`] (a `Vec<usize>` of arch ids
//! beside a zero-padded `Tensor<B, 2>` of weights) and inherent
//! `init`/`ask`/`tell`/`best` methods mirroring the `Strategy` shape — it does
//! **not** implement `Strategy<B>`.
//!
//! # Architecture dispatch — closure-erased registry
//!
//! Because [`ParamReshaper`] carries an
//! associated `type Module`, a `dyn ParamReshaper` is not object-safe and a
//! `Vec` of heterogeneous-variant reshapers is impossible. Instead each variant
//! is type-erased into a [`VariantEvaluator`]: a boxed closure that owns a
//! concrete [`ModuleReshaper<B, Vi>`](crate::param_reshaper::ModuleReshaper),
//! slices the active prefix of a (padded) flat weight row, unflattens it into
//! the concrete `Vi`, and scores it. Dispatch is by `arch_id` index. The
//! concrete variant types never appear in [`ArchNasStrategy`] or
//! [`ArchNasFitnessFn`] generics — both are generic only over `B`.
//!
//! # The alignment invariant
//!
//! The single load-bearing invariant is that `arch_id` indexes the **same**
//! variant in the strategy (which uses per-variant parameter counts to pad and
//! re-initialize weights) and in the fitness adapter (which holds the scoring
//! closures). [`ArchNasBuilder`] makes this structural: it is the *single*
//! registration point, and [`ArchNasBuilder::build`] emits both [`NasParams`]
//! and [`ArchNasFitnessFn`] from one ordered list, so the indices cannot drift.
//!
//! # Gradient isolation and host RNG
//!
//! Every type here is generic over `B: Backend`, never `AutodiffBackend` —
//! gradient isolation at the type level. All sampling
//! (architecture ids, selection, crossover, weight perturbation) goes through
//! [`seed_stream`] host-side `StdRng` substreams; the
//! process-wide backend RNG is never touched (see [`crate::rng`]).

use std::marker::PhantomData;

use burn::module::Module;
use burn::tensor::{Tensor, TensorData, backend::Backend};
use rand::{Rng, RngExt};
use rand_distr::{Distribution as _, Normal};

use rlevo_core::config::{self, ConfigError, ConstraintKind, Validate};

use crate::ops::selection::{argmax_host, tournament_indices_host, truncation_indices_host};
use crate::param_reshaper::{ModuleReshaper, ParamReshaper};
use crate::rng::{SeedPurpose, seed_stream};

/// One population's worth of architecture-NAS genomes.
///
/// `arch_ids[i]` is the architecture index of individual `i` (in
/// `0..num_variants`); row `i` of `weights` holds that individual's flat
/// weight vector, padded with trailing zeros out to `max_param_count`.
///
/// # Not `Clone`
///
/// `NasGenome` deliberately does **not** derive [`Clone`]: Burn tensors alias
/// their underlying storage, so a derived `Clone` would silently share buffers
/// rather than copy them. [`ArchNasStrategy::ask`] constructs a *fresh*
/// `NasGenome` with newly allocated tensors (built host-side via
/// [`Tensor::from_data`]) instead of cloning the struct — consistent with
/// [`WeightOnly`](super::weight_only::WeightOnly) and the phase-1 strategies.
#[derive(Debug)]
pub struct NasGenome<B: Backend> {
    /// Architecture index of each individual; length `pop_size`, each value in
    /// `0..num_variants`.
    arch_ids: Vec<usize>,
    /// Per-individual flat weights, shape `[pop_size, max_param_count]`. Columns
    /// beyond a variant's own parameter count are zero-padded.
    weights: Tensor<B, 2>,
}

impl<B: Backend> NasGenome<B> {
    /// Assembles a population genome, checking one architecture id per row.
    ///
    /// # Errors
    ///
    /// Returns a [`ConfigError`] if `weights` has zero rows or if `arch_ids`
    /// does not have exactly one entry per weight row (`pop_size`).
    pub fn try_new(arch_ids: Vec<usize>, weights: Tensor<B, 2>) -> Result<Self, ConfigError> {
        let pop = weights.dims()[0];
        config::nonzero("NasGenome", "pop_size", pop)?;
        if arch_ids.len() != pop {
            return Err(ConfigError {
                config: "NasGenome",
                field: "arch_ids",
                kind: ConstraintKind::Custom("must have one architecture id per weight row"),
            });
        }
        Ok(Self { arch_ids, weights })
    }

    /// Architecture index of each individual, length `pop_size`.
    #[must_use]
    pub fn arch_ids(&self) -> &[usize] {
        &self.arch_ids
    }

    /// Per-individual flat weights, shape `[pop_size, max_param_count]`.
    #[must_use]
    pub fn weights(&self) -> &Tensor<B, 2> {
        &self.weights
    }
}

/// A type-erased per-variant evaluator.
///
/// Built from a concrete `Module` variant `M` and a scorer `Fn(&M) -> f32`, it
/// captures a [`ModuleReshaper<B, M>`](crate::param_reshaper::ModuleReshaper)
/// in a boxed closure and exposes only `B`-generic operations. The concrete
/// `M` is invisible to callers.
pub struct VariantEvaluator<B: Backend> {
    num_params: usize,
    score_fn: Box<dyn Fn(Tensor<B, 1>) -> f32 + Send + Sync>,
}

impl<B: Backend> std::fmt::Debug for VariantEvaluator<B> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VariantEvaluator")
            .field("num_params", &self.num_params)
            .finish_non_exhaustive()
    }
}

impl<B: Backend> VariantEvaluator<B> {
    /// Build an evaluator for one architecture variant.
    ///
    /// `template` is an instance of the concrete `Module` variant (its weights
    /// are irrelevant — only its shape is used, to build the reshaper).
    /// `scorer` maps an unflattened module to a fitness value in the canonical
    /// **maximise** convention (higher is better).
    #[must_use]
    pub fn new<M, F>(template: M, scorer: F) -> Self
    where
        M: Module<B> + Sync + 'static,
        F: Fn(&M) -> f32 + Send + Sync + 'static,
    {
        let reshaper = ModuleReshaper::new(template);
        let num_params = reshaper.num_params();
        let score_fn = Box::new(move |active: Tensor<B, 1>| {
            let module = reshaper.unflatten(active);
            scorer(&module)
        });
        Self {
            num_params,
            score_fn,
        }
    }

    /// Number of flat float parameters this variant occupies (the active
    /// prefix width of a padded weight row).
    #[must_use]
    pub fn num_params(&self) -> usize {
        self.num_params
    }

    /// Score a single (padded) weight row.
    ///
    /// The leading `num_params` columns are sliced off, unflattened into the
    /// concrete module, and scored; trailing padding is ignored. Returns a
    /// fitness value in the canonical maximise convention.
    ///
    /// # Panics
    ///
    /// Panics if `padded_row` has fewer than `num_params` elements.
    #[must_use]
    pub fn score(&self, padded_row: Tensor<B, 1>) -> f32 {
        #[allow(clippy::single_range_in_vec_init)]
        let active = padded_row.slice([0..self.num_params]);
        (self.score_fn)(active)
    }
}

/// Static configuration for an [`ArchNasStrategy`] run. No `B` generic.
///
/// Produced by [`ArchNasBuilder::build`] so `per_variant_params` and
/// `max_param_count` are derived from the registered variants and cannot drift
/// from the fitness adapter's registry.
#[derive(Debug, Clone)]
pub struct NasParams {
    /// Number of individuals per generation.
    pop_size: usize,
    /// Number of architecture variants (`arch_id` ranges over `0..num_variants`).
    num_variants: usize,
    /// Flat parameter count of each variant; length `num_variants`, indexed by
    /// `arch_id`.
    per_variant_params: Vec<usize>,
    /// Width of the weight tensor — `per_variant_params.iter().max()`.
    max_param_count: usize,
    /// Per-individual probability of reassigning the architecture id each
    /// generation (a changed id re-initializes that child's active weights).
    arch_mutation_rate: f32,
    /// Standard deviation of the isotropic Gaussian weight perturbation.
    weight_mutation_std: f32,
    /// Standard deviation used to initialize (and re-initialize) weights.
    weight_init_std: f32,
    /// Tournament size for parent selection (clamped to `>= 1`).
    tournament_size: usize,
    /// Number of best individuals carried over unmutated each generation
    /// (clamped to `<= pop_size`).
    elite_count: usize,
}

impl NasParams {
    /// Number of individuals per generation.
    #[must_use]
    pub fn pop_size(&self) -> usize {
        self.pop_size
    }

    /// Number of architecture variants (`arch_id` ranges over `0..num_variants`).
    #[must_use]
    pub fn num_variants(&self) -> usize {
        self.num_variants
    }

    /// Flat parameter count of each variant, indexed by `arch_id`.
    #[must_use]
    pub fn per_variant_params(&self) -> &[usize] {
        &self.per_variant_params
    }

    /// Width of the weight tensor (`per_variant_params.iter().max()`).
    #[must_use]
    pub fn max_param_count(&self) -> usize {
        self.max_param_count
    }

    /// Per-individual probability of reassigning the architecture id each
    /// generation.
    #[must_use]
    pub fn arch_mutation_rate(&self) -> f32 {
        self.arch_mutation_rate
    }

    /// Standard deviation of the isotropic Gaussian weight perturbation.
    #[must_use]
    pub fn weight_mutation_std(&self) -> f32 {
        self.weight_mutation_std
    }

    /// Standard deviation used to initialize (and re-initialize) weights.
    #[must_use]
    pub fn weight_init_std(&self) -> f32 {
        self.weight_init_std
    }

    /// Tournament size for parent selection.
    #[must_use]
    pub fn tournament_size(&self) -> usize {
        self.tournament_size
    }

    /// Number of best individuals carried over unmutated each generation.
    #[must_use]
    pub fn elite_count(&self) -> usize {
        self.elite_count
    }
}

impl Validate for NasParams {
    fn validate(&self) -> Result<(), ConfigError> {
        const C: &str = "NasParams";
        config::at_least(C, "pop_size", self.pop_size, 1)?;
        config::at_least(C, "num_variants", self.num_variants, 1)?;
        if self.per_variant_params.len() != self.num_variants {
            return Err(ConfigError {
                config: C,
                field: "per_variant_params",
                kind: ConstraintKind::Custom("per_variant_params length must equal num_variants"),
            });
        }
        config::at_least(C, "tournament_size", self.tournament_size, 1)?;
        if self.elite_count > self.pop_size {
            return Err(ConfigError {
                config: C,
                field: "elite_count",
                kind: ConstraintKind::Custom("elite_count must not exceed pop_size"),
            });
        }
        Ok(())
    }
}

/// Evolving state for [`ArchNasStrategy`]: resident population plus best-ever
/// tracking.
///
/// Carries no PRNG state (the harness owns all stochasticity). Not `Clone`
/// (it holds a [`NasGenome`], which is intentionally not `Clone`); the strategy
/// threads state by explicit reconstruction.
#[derive(Debug)]
pub struct NasState<B: Backend> {
    population: NasGenome<B>,
    /// Resident fitness; empty until the first [`tell`](ArchNasStrategy::tell).
    fitness: Vec<f32>,
    best_arch_id: Option<usize>,
    /// Best-ever padded weight row, length `max_param_count`.
    best_weights: Option<Tensor<B, 1>>,
    best_fitness: f32,
    generation: usize,
}

impl<B: Backend> NasState<B> {
    /// Rebuild this state with freshly cloned tensors so it can be handed from
    /// `ask` to `tell` without sharing the resident buffers.
    fn carry_forward(&self) -> Self {
        Self {
            population: NasGenome {
                arch_ids: self.population.arch_ids.clone(),
                weights: self.population.weights.clone(),
            },
            fitness: self.fitness.clone(),
            best_arch_id: self.best_arch_id,
            best_weights: self.best_weights.clone(),
            best_fitness: self.best_fitness,
            generation: self.generation,
        }
    }

    /// Completed-generation counter (number of `tell` calls).
    #[must_use]
    pub fn generation(&self) -> usize {
        self.generation
    }

    /// Borrow the current resident population.
    ///
    /// After [`init`](ArchNasStrategy::init) this is the seed population; after
    /// each [`tell`](ArchNasStrategy::tell) it is the most recently evaluated
    /// population. Exposed so callers can inspect architecture coverage and the
    /// padded weight tensor.
    #[must_use]
    pub fn population(&self) -> &NasGenome<B> {
        &self.population
    }

    /// Best-ever fitness seen so far (canonical maximise), or `−∞` before the
    /// first [`tell`](ArchNasStrategy::tell).
    #[must_use]
    pub fn best_fitness(&self) -> f32 {
        self.best_fitness
    }
}

/// Knobs [`ArchNasBuilder::build`] cannot derive from the registered variants.
#[derive(Debug, Clone, Copy)]
pub struct NasBuilderConfig {
    /// Number of individuals per generation.
    pub pop_size: usize,
    /// Per-individual architecture-mutation probability.
    pub arch_mutation_rate: f32,
    /// Gaussian weight-perturbation standard deviation.
    pub weight_mutation_std: f32,
    /// Gaussian weight-initialization standard deviation.
    pub weight_init_std: f32,
    /// Tournament size for parent selection.
    pub tournament_size: usize,
    /// Number of elites preserved unmutated each generation.
    pub elite_count: usize,
}

impl Validate for NasBuilderConfig {
    fn validate(&self) -> Result<(), ConfigError> {
        const C: &str = "NasBuilderConfig";
        config::at_least(C, "pop_size", self.pop_size, 1)?;
        config::in_range(
            C,
            "arch_mutation_rate",
            0.0,
            1.0,
            f64::from(self.arch_mutation_rate),
        )?;
        config::in_range(
            C,
            "weight_mutation_std",
            0.0,
            f64::INFINITY,
            f64::from(self.weight_mutation_std),
        )?;
        config::in_range(
            C,
            "weight_init_std",
            0.0,
            f64::INFINITY,
            f64::from(self.weight_init_std),
        )?;
        config::at_least(C, "tournament_size", self.tournament_size, 1)?;
        if self.elite_count > self.pop_size {
            return Err(ConfigError {
                config: C,
                field: "elite_count",
                kind: ConstraintKind::Custom("elite_count must not exceed pop_size"),
            });
        }
        Ok(())
    }
}

/// Builder that registers architecture variants in order and emits a matched
/// `(NasParams, ArchNasFitnessFn)` pair.
///
/// `add_variant` is the single registration point, so `arch_id == registration
/// index` for both the strategy and the fitness adapter — the alignment
/// invariant is enforced structurally, not by convention.
///
/// # Example
///
/// ```ignore
/// let mut builder = ArchNasBuilder::<B>::new();
/// builder
///     .add_variant(ShallowMlp::new(&device), shallow_scorer)
///     .add_variant(DeepMlp::new(&device), deep_scorer);
/// let (params, fitness) = builder.build(NasBuilderConfig { /* … */ });
/// ```
#[derive(Debug, Default)]
pub struct ArchNasBuilder<B: Backend> {
    evaluators: Vec<VariantEvaluator<B>>,
}

impl<B: Backend> ArchNasBuilder<B> {
    /// Create an empty builder.
    #[must_use]
    pub fn new() -> Self {
        Self {
            evaluators: Vec::new(),
        }
    }

    /// Register one architecture variant. The registration order is the
    /// `arch_id` assigned to this variant.
    pub fn add_variant<M, F>(&mut self, template: M, scorer: F) -> &mut Self
    where
        M: Module<B> + Sync + 'static,
        F: Fn(&M) -> f32 + Send + Sync + 'static,
    {
        self.evaluators
            .push(VariantEvaluator::new(template, scorer));
        self
    }

    /// Consume the builder, producing a matched configuration and fitness
    /// adapter from the same ordered variant list.
    ///
    /// # Panics
    ///
    /// Panics if no variants were registered.
    #[must_use]
    pub fn build(self, cfg: NasBuilderConfig) -> (NasParams, ArchNasFitnessFn<B>) {
        assert!(
            !self.evaluators.is_empty(),
            "ArchNasBuilder requires at least one registered variant"
        );
        let per_variant_params: Vec<usize> = self
            .evaluators
            .iter()
            .map(VariantEvaluator::num_params)
            .collect();
        let max_param_count = per_variant_params
            .iter()
            .copied()
            .max()
            .expect("non-empty variants");
        let params = NasParams {
            pop_size: cfg.pop_size,
            num_variants: self.evaluators.len(),
            per_variant_params,
            max_param_count,
            arch_mutation_rate: cfg.arch_mutation_rate,
            weight_mutation_std: cfg.weight_mutation_std,
            weight_init_std: cfg.weight_init_std,
            tournament_size: cfg.tournament_size,
            elite_count: cfg.elite_count,
        };
        let fitness = ArchNasFitnessFn {
            evaluators: self.evaluators,
        };
        (params, fitness)
    }
}

/// Arch-dispatched, loop-over-N fitness adapter; owns the type-erased registry.
///
/// `evaluate` walks the population row by row, dispatches each to its variant's
/// evaluator by `arch_id` index, and assembles a `[pop_size]` fitness tensor.
/// Burn 0.21 has no batched/`vmap` forward, so evaluation is loop-over-N; a
/// batched arch-dispatched path is a future addition.
#[derive(Debug)]
pub struct ArchNasFitnessFn<B: Backend> {
    evaluators: Vec<VariantEvaluator<B>>,
}

impl<B: Backend> ArchNasFitnessFn<B> {
    /// Number of registered architecture variants.
    #[must_use]
    pub fn num_variants(&self) -> usize {
        self.evaluators.len()
    }

    /// Evaluate a population, returning `[pop_size]` fitness (canonical maximise).
    ///
    /// Row `i` is dispatched to `evaluators[genome.arch_ids[i]]`.
    ///
    /// # Panics
    ///
    /// Panics if any `arch_id` is out of range for the registry, or if the
    /// weight tensor's row count disagrees with `arch_ids.len()`.
    #[must_use]
    pub fn evaluate(&self, genome: &NasGenome<B>, device: &B::Device) -> Tensor<B, 1> {
        let [pop_size, max_param_count] = genome.weights.dims();
        assert_eq!(
            pop_size,
            genome.arch_ids.len(),
            "weights row count must equal arch_ids length"
        );
        let mut fitness: Vec<f32> = Vec::with_capacity(pop_size);
        for (i, &arch_id) in genome.arch_ids.iter().enumerate() {
            #[allow(clippy::single_range_in_vec_init)]
            let row: Tensor<B, 1> = genome
                .weights
                .clone()
                .slice([i..i + 1])
                .reshape([max_param_count]);
            fitness.push(self.evaluators[arch_id].score(row));
        }
        Tensor::<B, 1>::from_data(TensorData::new(fitness, [pop_size]), device)
    }
}

/// Custom NAS harness. Does **not** implement
/// [`Strategy`](crate::strategy::Strategy) — see the module docs.
///
/// The strategy is stateless; all evolving state lives in [`NasState`]. Drive
/// it with the same ask/tell loop as a [`Strategy`](crate::strategy::Strategy):
///
/// ```ignore
/// let strat = ArchNasStrategy::<B>::new();
/// let mut state = strat.init(&params, &mut rng, &device);
/// for _ in 0..generations {
///     let (genome, next) = strat.ask(&params, &state, &mut rng, &device);
///     let fitness = fitness_fn.evaluate(&genome, &device);
///     state = strat.tell(&params, genome, fitness, next, &mut rng);
/// }
/// let (arch_id, weights, fitness) = strat.best(&state).unwrap();
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct ArchNasStrategy<B: Backend> {
    _backend: PhantomData<fn() -> B>,
}

impl<B: Backend> ArchNasStrategy<B> {
    /// Build a new (stateless) NAS strategy.
    #[must_use]
    pub fn new() -> Self {
        Self {
            _backend: PhantomData,
        }
    }

    /// Build the initial state: uniform random architecture ids and
    /// Gaussian-initialized active weights (padding zeroed).
    ///
    /// `fitness` is left empty and `best_fitness` is `−∞`; the first
    /// [`tell`](Self::tell) populates both.
    ///
    /// # Panics
    ///
    /// Panics if `weight_init_std` is non-finite (`+∞` or `NaN`).
    #[must_use]
    pub fn init(
        &self,
        params: &NasParams,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> NasState<B> {
        debug_assert!(
            params.validate().is_ok(),
            "invalid NasParams reached init: {params:?}"
        );
        let pop = params.pop_size;
        let max = params.max_param_count;

        let mut arch_rng = seed_stream(rng.next_u64(), 0, SeedPurpose::Representative);
        let arch_ids: Vec<usize> = (0..pop)
            .map(|_| arch_rng.random_range(0..params.num_variants))
            .collect();

        let mut weight_rng = seed_stream(rng.next_u64(), 0, SeedPurpose::Init);
        let normal = Normal::new(0.0f32, params.weight_init_std).unwrap_or_else(|err| {
            panic!(
                "weight_init_std must be finite, got {}: {err}",
                params.weight_init_std
            )
        });
        let mut data = vec![0.0f32; pop * max];
        for (i, &arch) in arch_ids.iter().enumerate() {
            let n = params.per_variant_params[arch];
            for slot in &mut data[i * max..i * max + n] {
                *slot = normal.sample(&mut weight_rng);
            }
        }
        let weights = Tensor::<B, 2>::from_data(TensorData::new(data, [pop, max]), device);

        NasState {
            population: NasGenome { arch_ids, weights },
            fitness: Vec::new(),
            best_arch_id: None,
            best_weights: None,
            best_fitness: f32::NEG_INFINITY,
            generation: 0,
        }
    }

    /// Propose the next population.
    ///
    /// Before the first [`tell`](Self::tell) (resident fitness still empty) the
    /// unchanged resident population is returned for evaluation. Afterwards the
    /// method produces offspring: `elite_count` best residents carried over
    /// unmutated, then tournament selection, same-architecture blend crossover
    /// (single-parent copy when parents' architectures differ), and either a
    /// per-child architecture mutation (re-initializing the child's active
    /// weights) at `arch_mutation_rate` or an isotropic Gaussian weight
    /// perturbation. Five independent host substreams are derived from `rng`:
    /// selection, crossover, mutation, architecture (re)assignment, and weight
    /// re-initialization.
    ///
    /// The returned [`NasGenome`] is freshly allocated (no struct-level clone).
    ///
    /// # Panics
    ///
    /// Panics if `weight_init_std` or `weight_mutation_std` is non-finite
    /// (`+∞` or `NaN`), or if the resident weight tensor cannot be read back to
    /// the host as `f32`.
    #[must_use]
    pub fn ask(
        &self,
        params: &NasParams,
        state: &NasState<B>,
        rng: &mut dyn Rng,
        device: &<B as burn::tensor::backend::BackendTypes>::Device,
    ) -> (NasGenome<B>, NasState<B>) {
        let pop = params.pop_size;
        let max = params.max_param_count;

        // First call: harness has not evaluated the seed population yet.
        if state.fitness.is_empty() {
            let genome = NasGenome {
                arch_ids: state.population.arch_ids.clone(),
                weights: state.population.weights.clone(),
            };
            return (genome, state.carry_forward());
        }

        let gen_idx = state.generation as u64;
        let mut sel_rng = seed_stream(rng.next_u64(), gen_idx, SeedPurpose::Selection);
        let mut xover_rng = seed_stream(rng.next_u64(), gen_idx, SeedPurpose::Crossover);
        let mut mut_rng = seed_stream(rng.next_u64(), gen_idx, SeedPurpose::Mutation);
        let mut arch_rng = seed_stream(rng.next_u64(), gen_idx, SeedPurpose::Representative);
        let mut winit_rng = seed_stream(rng.next_u64(), gen_idx, SeedPurpose::Other);

        let perturb = Normal::new(0.0f32, params.weight_mutation_std).unwrap_or_else(|err| {
            panic!(
                "weight_mutation_std must be finite, got {}: {err}",
                params.weight_mutation_std
            )
        });
        let reinit = Normal::new(0.0f32, params.weight_init_std).unwrap_or_else(|err| {
            panic!(
                "weight_init_std must be finite, got {}: {err}",
                params.weight_init_std
            )
        });

        let resident: Vec<f32> = state
            .population
            .weights
            .clone()
            .into_data()
            .into_vec::<f32>()
            .expect("weights tensor must be readable as f32");
        let resident_arch = &state.population.arch_ids;

        // Elitism: indices sorted by descending (better, canonical maximise)
        // fitness — highest first; NaN sanitised to −inf so it never ranks
        // as best.
        let order: Vec<usize> = truncation_indices_host(&state.fitness, pop)
            .into_iter()
            .map(|i| usize::try_from(i).expect("winner index is non-negative"))
            .collect();
        let elite_count = params.elite_count.min(pop);
        let tournament_size = params.tournament_size.max(1);

        let mut child_arch: Vec<usize> = Vec::with_capacity(pop);
        let mut child: Vec<f32> = vec![0.0f32; pop * max];

        // Carry elites over unmutated.
        for (ci, &ei) in order[..elite_count].iter().enumerate() {
            child[ci * max..ci * max + max].copy_from_slice(&resident[ei * max..ei * max + max]);
            child_arch.push(resident_arch[ei]);
        }

        // Fill the rest with offspring. All parent tournaments draw from
        // the dedicated `sel_rng` stream, so batching the draws up front
        // preserves the exact per-child draw order.
        let parents = tournament_indices_host(
            &state.fitness,
            tournament_size,
            2 * (pop - elite_count),
            &mut sel_rng,
        );
        for ci in elite_count..pop {
            let pair = 2 * (ci - elite_count);
            let pa = usize::try_from(parents[pair]).expect("winner index is non-negative");
            let pb = usize::try_from(parents[pair + 1]).expect("winner index is non-negative");
            let arch = resident_arch[pa];
            let n = params.per_variant_params[arch];
            let base = ci * max;

            if resident_arch[pa] == resident_arch[pb] {
                // Same-architecture blend crossover on the active region.
                for j in 0..n {
                    let alpha: f32 = xover_rng.random::<f32>();
                    child[base + j] =
                        alpha * resident[pa * max + j] + (1.0 - alpha) * resident[pb * max + j];
                }
            } else {
                // Architecture mismatch: single-parent copy (mutation-only).
                child[base..base + n].copy_from_slice(&resident[pa * max..pa * max + n]);
            }

            if arch_rng.random::<f32>() < params.arch_mutation_rate {
                // Architecture mutation: switch variant, re-initialize weights.
                let new_arch = arch_rng.random_range(0..params.num_variants);
                let nn = params.per_variant_params[new_arch];
                child[base..base + max].fill(0.0);
                for slot in &mut child[base..base + nn] {
                    *slot = reinit.sample(&mut winit_rng);
                }
                child_arch.push(new_arch);
            } else {
                // Isotropic Gaussian weight perturbation on the active region.
                for slot in &mut child[base..base + n] {
                    *slot += perturb.sample(&mut mut_rng);
                }
                child_arch.push(arch);
            }
        }

        let weights = Tensor::<B, 2>::from_data(TensorData::new(child, [pop, max]), device);
        let genome = NasGenome {
            arch_ids: child_arch,
            weights,
        };
        (genome, state.carry_forward())
    }

    /// Consume a population's fitness and produce the next state.
    ///
    /// Records the told population as the new residents, updates the best-ever
    /// `(arch_id, weights, fitness)` triple, and increments the generation.
    /// `fitness` follows the canonical maximise convention (higher is better).
    ///
    /// # Fitness hygiene
    ///
    /// This is the `ArchNasStrategy` **driver chokepoint** (ADR 0034): the
    /// strategy is its own driver and does **not** run through
    /// [`EvolutionaryHarness`](crate::strategy::EvolutionaryHarness), so there is
    /// no harness above it. `tell` applies
    /// `sanitize_fitness` to the host fitness
    /// vector (`NaN → −∞`, `+∞ → f32::MAX`, `−∞` and finite pass through)
    /// **before** `update_best` and before it is stored as `state.fitness`, so
    /// a non-finite fitness can never become an immortal champion nor win the
    /// next [`ask`](Self::ask)'s elite carry or tournament.
    ///
    /// # Panics
    ///
    /// Panics if the `fitness` tensor cannot be read back to the host as `f32`.
    #[must_use]
    pub fn tell(
        &self,
        params: &NasParams,
        population: NasGenome<B>,
        fitness: Tensor<B, 1>,
        mut state: NasState<B>,
        _rng: &mut dyn Rng,
    ) -> NasState<B> {
        let raw = fitness
            .into_data()
            .into_vec::<f32>()
            .expect("fitness tensor must be readable as f32");
        // Driver chokepoint (ADR 0034): sanitize before update_best/store so a
        // NaN/±∞ can neither become the champion nor pollute the next ask.
        let fitness_host: Vec<f32> = raw
            .into_iter()
            .map(crate::fitness::sanitize_fitness)
            .collect();
        update_best(
            &mut state,
            &population,
            &fitness_host,
            params.max_param_count,
        );
        state.population = population;
        state.fitness = fitness_host;
        state.generation += 1;
        state
    }

    /// Best-ever individual: `(arch_id, padded weights, fitness)`.
    ///
    /// The weight tensor has length `max_param_count` (trailing padding
    /// retained). Returns `None` before the first [`tell`](Self::tell).
    #[must_use]
    pub fn best(&self, state: &NasState<B>) -> Option<(usize, Tensor<B, 1>, f32)> {
        match (state.best_arch_id, state.best_weights.as_ref()) {
            (Some(arch_id), Some(weights)) => Some((arch_id, weights.clone(), state.best_fitness)),
            _ => None,
        }
    }
}

/// Update best-ever tracking from a freshly-evaluated population.
///
/// Both the argmax ([`argmax_host`], seeded at `f32::NEG_INFINITY`) and the
/// running champion (`state.best_fitness`, initialised to `f32::NEG_INFINITY`
/// in [`init`](ArchNasStrategy::init)) start from the maximise-space worst
/// sentinel, so a leading broken member cannot suppress a later valid best.
/// `fitness` is expected pre-sanitized by [`tell`](ArchNasStrategy::tell).
fn update_best<B: Backend>(
    state: &mut NasState<B>,
    pop: &NasGenome<B>,
    fitness: &[f32],
    max_param_count: usize,
) {
    if fitness.is_empty() {
        return;
    }
    let best_idx = argmax_host(fitness);
    let best_f = fitness[best_idx];
    if best_f > state.best_fitness {
        #[allow(clippy::single_range_in_vec_init)]
        let row: Tensor<B, 1> = pop
            .weights
            .clone()
            .slice([best_idx..best_idx + 1])
            .reshape([max_param_count]);
        state.best_weights = Some(row);
        state.best_arch_id = Some(pop.arch_ids[best_idx]);
        state.best_fitness = best_f;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use burn::backend::Flex;
    use burn::nn::{Linear, LinearConfig};
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    type TestBackend = Flex;

    #[test]
    fn nas_genome_try_new_checks_one_arch_id_per_row() {
        let device = Default::default();
        let weights = Tensor::<TestBackend, 2>::zeros([3, 4], &device);
        assert!(NasGenome::try_new(vec![0, 1, 2], weights).is_ok());
        let weights = Tensor::<TestBackend, 2>::zeros([3, 4], &device);
        assert!(NasGenome::try_new(vec![0, 1], weights).is_err());
        let empty = Tensor::<TestBackend, 2>::zeros([0, 4], &device);
        assert!(NasGenome::try_new(vec![], empty).is_err());
    }

    fn valid_builder_config() -> NasBuilderConfig {
        NasBuilderConfig {
            pop_size: 16,
            arch_mutation_rate: 0.1,
            weight_mutation_std: 0.1,
            weight_init_std: 0.5,
            tournament_size: 2,
            elite_count: 1,
        }
    }

    #[test]
    fn builder_config_validates() {
        assert!(valid_builder_config().validate().is_ok());
    }

    #[test]
    fn builder_config_rejects_elite_above_pop() {
        let mut cfg = valid_builder_config();
        cfg.elite_count = 32;
        assert_eq!(cfg.validate().unwrap_err().field, "elite_count");
    }

    #[test]
    fn nas_params_validate_and_reject() {
        let good = NasParams {
            pop_size: 16,
            num_variants: 2,
            per_variant_params: vec![10, 20],
            max_param_count: 20,
            arch_mutation_rate: 0.1,
            weight_mutation_std: 0.1,
            weight_init_std: 0.5,
            tournament_size: 2,
            elite_count: 1,
        };
        assert!(good.validate().is_ok());
        let mut bad = good.clone();
        bad.per_variant_params = vec![10];
        assert_eq!(bad.validate().unwrap_err().field, "per_variant_params");
    }

    /// One-hidden-layer MLP variant: `2 -> H -> 1`.
    #[derive(Module, Debug)]
    struct ShallowMlp<B: Backend> {
        l1: Linear<B>,
        l2: Linear<B>,
    }

    impl<B: Backend> ShallowMlp<B> {
        fn new(hidden: usize, device: &B::Device) -> Self {
            Self {
                l1: LinearConfig::new(2, hidden).init(device),
                l2: LinearConfig::new(hidden, 1).init(device),
            }
        }
    }

    /// Two-hidden-layer MLP variant: `2 -> H -> H/2 -> 1`.
    #[derive(Module, Debug)]
    struct DeepMlp<B: Backend> {
        l1: Linear<B>,
        l2: Linear<B>,
        l3: Linear<B>,
    }

    impl<B: Backend> DeepMlp<B> {
        fn new(device: &B::Device) -> Self {
            Self {
                l1: LinearConfig::new(2, 8).init(device),
                l2: LinearConfig::new(8, 4).init(device),
                l3: LinearConfig::new(4, 1).init(device),
            }
        }
    }

    // `device` is borrowed to mirror the production `&B::Device` convention;
    // the Flex test device is zero-sized, hence the targeted allow.
    #[allow(clippy::trivially_copy_pass_by_ref)]
    fn two_variant_builder(
        device: &<TestBackend as burn::tensor::backend::BackendTypes>::Device,
    ) -> ArchNasBuilder<TestBackend> {
        // Two trivial scorers (constant 0.0) — these tests exercise only the
        // builder/registry plumbing, not fitness dynamics.
        fn zero_shallow(_m: &ShallowMlp<TestBackend>) -> f32 {
            0.0
        }
        fn zero_deep(_m: &DeepMlp<TestBackend>) -> f32 {
            0.0
        }
        let mut builder = ArchNasBuilder::<TestBackend>::new();
        builder
            // arch 0: shallow (2*4 + 4 + 4*1 + 1 = 17 params)
            .add_variant(ShallowMlp::<TestBackend>::new(4, device), zero_shallow)
            // arch 1: deep (2*8 + 8 + 8*4 + 4 + 4*1 + 1 = 65 params)
            .add_variant(DeepMlp::<TestBackend>::new(device), zero_deep);
        builder
    }

    #[test]
    fn builder_aligns_arch_id_with_param_counts() {
        let device = Default::default();
        let (params, fitness) = two_variant_builder(&device).build(NasBuilderConfig {
            pop_size: 8,
            arch_mutation_rate: 0.1,
            weight_mutation_std: 0.1,
            weight_init_std: 0.5,
            tournament_size: 2,
            elite_count: 1,
        });
        assert_eq!(params.num_variants, 2);
        assert_eq!(params.per_variant_params, vec![17, 65]);
        assert_eq!(params.max_param_count, 65);
        assert_eq!(fitness.num_variants(), 2);
    }

    #[test]
    fn variant_evaluator_dispatches_and_slices_active_prefix() {
        let device = Default::default();
        // Scorer reports the number of params it sees, proving the active
        // prefix (not the padded width) reaches the concrete module.
        let mut builder = ArchNasBuilder::<TestBackend>::new();
        builder.add_variant(
            ShallowMlp::<TestBackend>::new(4, &device),
            |m: &ShallowMlp<TestBackend>| {
                #[allow(clippy::cast_precision_loss)]
                let n = ModuleReshaper::new(ShallowMlp::<TestBackend>::new(4, &Default::default()))
                    .num_params() as f32;
                // forward to make the module non-dead; result unused beyond shape.
                let _ = &m.l1;
                n
            },
        );
        let (params, fitness) = builder.build(NasBuilderConfig {
            pop_size: 1,
            arch_mutation_rate: 0.0,
            weight_mutation_std: 0.0,
            weight_init_std: 0.1,
            tournament_size: 1,
            elite_count: 0,
        });
        // Build a one-row genome with arch 0 and 17 active params, padded to 17.
        let weights = Tensor::<TestBackend, 2>::from_data(
            TensorData::new(
                vec![0.0f32; params.max_param_count],
                [1, params.max_param_count],
            ),
            &device,
        );
        let genome = NasGenome {
            arch_ids: vec![0],
            weights,
        };
        let fit = fitness
            .evaluate(&genome, &device)
            .into_data()
            .into_vec::<f32>()
            .expect("fitness host-read of a tensor this test just built");
        approx::assert_relative_eq!(fit[0], 17.0, epsilon = 1e-6);
    }

    #[test]
    fn init_populates_all_variants_and_zero_pads() {
        let device = Default::default();
        let (params, _fitness) = two_variant_builder(&device).build(NasBuilderConfig {
            pop_size: 60,
            arch_mutation_rate: 0.1,
            weight_mutation_std: 0.1,
            weight_init_std: 0.5,
            tournament_size: 2,
            elite_count: 1,
        });
        let strat = ArchNasStrategy::<TestBackend>::new();
        let mut rng = StdRng::seed_from_u64(7);
        let state = strat.init(&params, &mut rng, &device);

        // Both architectures should be represented with pop 60.
        assert!(state.population.arch_ids.contains(&0));
        assert!(state.population.arch_ids.contains(&1));

        // A shallow (arch 0) row must have zeros beyond its 17 active params.
        let rows = state
            .population
            .weights
            .clone()
            .into_data()
            .into_vec::<f32>()
            .expect("population host-read of a tensor this test just built");
        let max = params.max_param_count;
        let shallow_row = state
            .population
            .arch_ids
            .iter()
            .position(|&a| a == 0)
            .unwrap();
        for &v in &rows[shallow_row * max + 17..shallow_row * max + max] {
            approx::assert_relative_eq!(v, 0.0, epsilon = 1e-6);
        }
    }

    #[test]
    fn ask_tell_runs_without_panic_and_tracks_best() {
        let device = Default::default();
        // Scorer: sum of squared active weights — a canonical (maximise)
        // fitness the strategy drives upward; the test only pins best-ever
        // monotonicity, not optimisation quality.
        let mut builder = ArchNasBuilder::<TestBackend>::new();
        let sq = |m: &ShallowMlp<TestBackend>| {
            let r = ModuleReshaper::new(ShallowMlp::<TestBackend>::new(4, &Default::default()));
            let flat = r.flatten(m, &Default::default());
            flat.clone()
                .mul(flat)
                .sum()
                .into_data()
                .into_vec::<f32>()
                .expect("output host-read of a tensor this test just built")[0]
        };
        let sq_deep = |m: &DeepMlp<TestBackend>| {
            let r = ModuleReshaper::new(DeepMlp::<TestBackend>::new(&Default::default()));
            let flat = r.flatten(m, &Default::default());
            flat.clone()
                .mul(flat)
                .sum()
                .into_data()
                .into_vec::<f32>()
                .expect("output host-read of a tensor this test just built")[0]
        };
        builder
            .add_variant(ShallowMlp::<TestBackend>::new(4, &device), sq)
            .add_variant(DeepMlp::<TestBackend>::new(&device), sq_deep);
        let (params, fitness) = builder.build(NasBuilderConfig {
            pop_size: 30,
            arch_mutation_rate: 0.1,
            weight_mutation_std: 0.05,
            weight_init_std: 0.5,
            tournament_size: 3,
            elite_count: 2,
        });

        let strat = ArchNasStrategy::<TestBackend>::new();
        let mut rng = StdRng::seed_from_u64(123);
        let mut state = strat.init(&params, &mut rng, &device);

        // Generation 0.
        let (genome, next) = strat.ask(&params, &state, &mut rng, &device);
        let fit = fitness.evaluate(&genome, &device);
        state = strat.tell(&params, genome, fit, next, &mut rng);
        let gen0_best = strat.best(&state).map(|(_, _, f)| f).unwrap();

        // Generations 1..6.
        for _ in 0..6 {
            let (genome, next) = strat.ask(&params, &state, &mut rng, &device);
            let fit = fitness.evaluate(&genome, &device);
            state = strat.tell(&params, genome, fit, next, &mut rng);
        }
        let final_best = strat.best(&state).map(|(_, _, f)| f).unwrap();

        assert!(
            final_best >= gen0_best,
            "best-ever must be monotone (maximise): final {final_best} < gen0 {gen0_best}"
        );
        assert_eq!(state.generation(), 7);
    }

    /// Regression (ADR 0034, issue #133): `tell` is the `ArchNasStrategy` chokepoint
    /// and must sanitize non-finite fitness before `update_best`/store. A `NaN`
    /// member must never become the champion nor survive into the next ask's
    /// selection; a `+∞` member must rank top but as a **finite** value
    /// (`f32::MAX`).
    #[test]
    fn tell_sanitizes_nan_and_inf_so_nan_never_champions() {
        let device = Default::default();
        let (params, _fitness) = two_variant_builder(&device).build(NasBuilderConfig {
            pop_size: 4,
            arch_mutation_rate: 0.0,
            weight_mutation_std: 0.0,
            weight_init_std: 0.5,
            tournament_size: 2,
            elite_count: 1,
        });
        let strat = ArchNasStrategy::<TestBackend>::new();
        let mut rng = StdRng::seed_from_u64(9);
        let state = strat.init(&params, &mut rng, &device);

        // One ask, then hand-craft the fitness: NaN (must never champion), +∞
        // (ranks top but finite), two plain values. Row 1 (+∞) is the champion.
        let (genome, next) = strat.ask(&params, &state, &mut rng, &device);
        let champion_arch = genome.arch_ids[1];
        let fitness = Tensor::<TestBackend, 1>::from_data(
            TensorData::new(vec![f32::NAN, f32::INFINITY, 1.0_f32, 2.0], [4]),
            &device,
        );
        let state = strat.tell(&params, genome, fitness, next, &mut rng);

        // Stored fitness is sanitized: no NaN; NaN → −∞, +∞ → f32::MAX.
        assert!(
            state.fitness.iter().all(|f| !f.is_nan()),
            "no stored fitness is NaN after the tell chokepoint"
        );
        assert!(
            state.fitness[0].is_infinite() && state.fitness[0].is_sign_negative(),
            "the NaN member is sanitized to the maximise-space worst sentinel (−∞)"
        );
        approx::assert_relative_eq!(state.fitness[1], f32::MAX);

        // The champion is the sanitized +∞ = f32::MAX (finite), at row 1 — not
        // the NaN row.
        let (best_arch, _weights, best_f) = strat.best(&state).expect("best exists after tell");
        assert!(
            best_f.is_finite(),
            "champion fitness is finite (never NaN/±∞); got {best_f}"
        );
        approx::assert_relative_eq!(best_f, f32::MAX);
        assert_eq!(
            best_arch, champion_arch,
            "champion is the +∞ row, never the NaN row"
        );
    }
}