Neural Ensemble Search (NAS)

def search(
    self,
    val_loader: DataLoader,
    val_loader_shift: DataLoader | None = None,
) -> Ensemble:
    """Run NES-RS and return the selected ensemble.

    Args:
        val_loader: Validation loader used for ensemble selection.
        val_loader_shift: If provided, used instead of `val_loader` for
            the final ForwardSelect call (dataset-shift adaptation).

    Returns:
        Ensemble of `ensemble_size` members.
    """
    pool: list[nn.Module] = []
    for _ in range(self.pool_size):
        config = self.space.sample()
        model = self.space.build(config)
        self.train_fn(model)
        pool.append(model)

    selection_loader = (
        val_loader_shift if val_loader_shift is not None else val_loader
    )
    selected = forward_select(
        pool, selection_loader, self.ensemble_size, self.device, self.criterion
    )
    return Ensemble.from_models(selected)

def search(
    self,
    val_loader: DataLoader,
    val_loader_shift: DataLoader | None = None,
) -> Ensemble:
    """Run NES-RE and return the selected ensemble.

    Args:
        val_loader: Validation loader used for parent selection and final
            ensemble selection.
        val_loader_shift: If provided, used instead of `val_loader` for
            the final ForwardSelect call (dataset-shift adaptation).

    Returns:
        Ensemble of `ensemble_size` members.
    """
    # Maps id(model) -> config to support mutation of selected parents.
    config_map: dict[int, dict] = {}

    def _build_and_train(config: dict) -> nn.Module:
        model = self.space.build(config)
        self.train_fn(model)
        config_map[id(model)] = config
        return model

    # --- Initialisation: seed population and pool ---
    population: deque[nn.Module] = deque()
    pool: list[nn.Module] = []

    for _ in range(self.population_size):
        model = _build_and_train(self.space.sample())
        population.append(model)
        pool.append(model)

    # --- Evolution loop ---
    while len(pool) < self.pool_size:
        # Select m parent candidates from the current population via ForwardSelect.
        parent_candidates = forward_select(
            list(population),
            val_loader,
            self.num_parent_candidates,
            self.device,
            self.criterion,
        )

        # Sample one parent uniformly at random.
        parent = random.choice(parent_candidates)
        parent_config = config_map[id(parent)]

        # Mutate and train child.
        child_config = self.space.mutate(parent_config)
        child = _build_and_train(child_config)

        population.append(child)
        pool.append(child)

        # Remove the oldest member from the population (regularized evolution).
        population.popleft()

    # --- Final ensemble selection ---
    selection_loader = (
        val_loader_shift if val_loader_shift is not None else val_loader
    )
    selected = forward_select(
        pool, selection_loader, self.ensemble_size, self.device, self.criterion
    )
    return Ensemble.from_models(selected)

def sample_mc(self, val_loader: DataLoader) -> Ensemble:
    """
    Args:
        val_loader (DataLoader): Used to evaluate the posterior.

    Returns:
        Ensemble: The final ensemble wrapped in bensemble's core abstraction.
    """

    candidates = self._build_pool(val_loader)
    probs = self._posterior_probs(candidates)
    chosen = torch.multinomial(
        probs, num_samples=self.ensemble_size, replacement=False
    )
    models = [candidates[idx].model for idx in chosen.tolist()]
    return Ensemble.from_models(models)

def sample_svgd(self, val_loader: DataLoader) -> Ensemble:
    """
    Args:
        val_loader (DataLoader): Used to evaluate the architecture's loss/posterior.

    Returns:
        Ensemble: The final ensemble wrapped in bensemble's core abstraction.
    """
    candidates = self._build_pool(val_loader)
    posterior = self._posterior_probs(candidates)
    n = len(candidates)
    logits = torch.log(posterior + 1e-8)

    particles = torch.multinomial(
        posterior, num_samples=self.ensemble_size, replacement=False
    )
    for _ in range(self.svgd_steps):
        for i in range(self.ensemble_size):
            current = particles[i].item()
            best_idx = current
            best_value = float("-inf")

            for candidate_idx in range(n):
                if candidate_idx in particles.tolist() and candidate_idx != current:
                    continue
                repulsion = 0.0
                for j in range(self.ensemble_size):
                    if j == i:
                        continue
                    other_idx = particles[j].item()
                    div = self._pairwise_diversity(
                        candidates[candidate_idx].probs,
                        candidates[other_idx].probs,
                    )
                    repulsion += div
                value = logits[candidate_idx].item() + (
                    self.svgd_lr * self.diversity_weight * repulsion
                )
                if value > best_value:
                    best_value = value
                    best_idx = candidate_idx
            particles[i] = best_idx

    models = [candidates[idx].model for idx in particles.tolist()]
    return Ensemble.from_models(models)

@abc.abstractmethod
def build(self, config: dict) -> nn.Module:
    """Instantiate and return an untrained nn.Module for the given config."""
    ...

@abc.abstractmethod
def mutate(self, config: dict) -> dict:
    """Return a new config that is a mutated copy of the given config.
    Must not modify the input config in place."""
    ...

@abc.abstractmethod
def sample(self) -> dict:
    """Sample a random architecture config uniformly from the space."""
    ...