path: root/sdil/conv_crossover.py

"""Matched non-backprop adapters for the audited CIFAR ResNet topology.

These adapters reuse :class:`CIFARLocalResNet`'s forward tensors, BatchNorm,
option-A shortcuts, and local optimizer.  They are kept separate from the
established SDIL/KP implementation so a crossover baseline cannot silently
change the already confirmed forward model.
"""
import torch
import torch.nn.functional as F

from .conv import CIFARLocalResNet


class CIFARPEPITAResNet(CIFARLocalResNet):
    """Two-presentation PEPITA on the matched residual forward topology."""

    def __init__(self, *args, projection_scale=0.05,
                 projection_seed=1731, **kwargs):
        super().__init__(*args, **kwargs)
        if projection_scale <= 0:
            raise ValueError("PEPITA projection scale must be positive")
        generator = torch.Generator(device="cpu").manual_seed(projection_seed)
        input_units = 3 * 32 * 32
        limit = (6.0 / input_units) ** 0.5 * projection_scale
        projection = (
            2.0 * torch.rand(
                self.n_classes, 3, 32, 32, generator=generator) - 1.0
        ) * limit
        self.input_feedback = projection.to(
            device=self.device, dtype=self.dtype)

    @property
    def n_fixed_feedback_parameters(self):
        return self.input_feedback.numel()

    @torch.no_grad()
    def pepita_ascent_directions(self, clean, modulated, modulated_error):
        """Return the explicit first-minus-second PEPITA correlations.

        The post-activation difference directly multiplies the modulated
        presynaptic activity, as in PEPITA/ERIN; it is not differentiated
        through the ReLU.  BatchNorm's current-layer Jacobian and a residual
        branch's fixed multiplier remain part of that local synaptic
        eligibility.  Convolutional correlations follow the reference code's
        additional average over spatial positions.
        """
        batch = modulated_error.shape[0]
        directions = []
        gamma_directions = []
        beta_directions = []
        for index, (clean_hidden, modulated_hidden, cache, weight, spec) in (
                enumerate(zip(
                    clean["hiddens"], modulated["hiddens"],
                    modulated["caches"], self.W, self.layer_specs))):
            field = clean_hidden - modulated_hidden
            local_field = field * spec.branch_scale
            local_field, gamma_direction, beta_direction = (
                self._normalization_backward(
                    index, local_field, cache["normalization"]))
            spatial = local_field.shape[2] * local_field.shape[3]
            correlation = torch.nn.grad.conv2d_weight(
                cache["pre"], weight.shape, local_field,
                stride=spec.stride, padding=spec.padding)
            directions.append(-correlation / (batch * spatial))
            if gamma_direction is not None:
                gamma_directions.append(
                    -gamma_direction / (batch * spatial))
                beta_directions.append(
                    -beta_direction / (batch * spatial))
        output_weight = -(
            modulated_error.t() @ modulated["features"]) / batch
        output_bias = -modulated_error.mean(dim=0)
        return (
            directions, gamma_directions, beta_directions,
            output_weight, output_bias)

    def pepita_step(self, image, labels, eta, eta_output=None,
                    momentum=0.0, weight_decay=0.0):
        """One architecture-compatible PEPITA/ERIN local update."""
        one_hot = F.one_hot(labels, self.n_classes).to(image.dtype)
        with torch.no_grad():
            clean = self.forward(
                image, return_cache=True, training=True, update_stats=True)
            clean_error = torch.softmax(clean["logits"], dim=1) - one_hot
            input_error = torch.einsum(
                "bc,cijk->bijk", clean_error, self.input_feedback)
            modulated = self.forward(
                image + input_error, return_cache=True,
                training=True, update_stats=False)
            modulated_error = (
                torch.softmax(modulated["logits"], dim=1) - one_hot)
            (directions, gamma_directions, beta_directions,
             output_weight, output_bias) = self.pepita_ascent_directions(
                 clean, modulated, modulated_error)
            self.apply_ascent(
                directions, output_weight, output_bias, eta,
                eta_output=eta_output, momentum=momentum,
                weight_decay=weight_decay,
                gamma_directions=gamma_directions,
                beta_directions=beta_directions)
            loss = F.cross_entropy(clean["logits"], labels)
        return float(loss)


class CIFARForwardForwardResNet(CIFARLocalResNet):
    """Greedy supervised Forward-Forward on the residual parameter topology."""

    def __init__(self, *args, threshold=2.0, learning_rate=0.03,
                 score_from_layer=1, **kwargs):
        super().__init__(*args, **kwargs)
        if learning_rate <= 0 or threshold <= 0:
            raise ValueError("invalid Forward-Forward hyperparameters")
        self.ff_threshold = float(threshold)
        self.ff_score_from_layer = int(score_from_layer)
        self.ff_num_layers = self.n_hidden + 1
        if not 0 <= self.ff_score_from_layer < self.ff_num_layers:
            raise ValueError("Forward-Forward score range excludes all layers")
        self.ff_optimizers = []
        for index, weight in enumerate(self.W):
            parameters = [weight]
            if self.normalization == "batchnorm":
                parameters.extend([self.gamma[index], self.beta[index]])
            for parameter in parameters:
                parameter.requires_grad_(True)
            self.ff_optimizers.append(
                torch.optim.Adam(parameters, lr=learning_rate))
        self.W_out.requires_grad_(True)
        self.b_out.requires_grad_(True)
        self.ff_optimizers.append(torch.optim.Adam(
            [self.W_out, self.b_out], lr=learning_rate))

    @staticmethod
    def _ff_normalize(value):
        axes = tuple(range(1, value.ndim))
        norm = torch.sqrt(torch.sum(value.square(), dim=axes, keepdim=True))
        return value / (norm + 1e-8)

    def ff_overlay(self, image, labels):
        flat = image.reshape(image.shape[0], -1).clone()
        flat[:, :self.n_classes] = 0.0
        flat[torch.arange(image.shape[0], device=image.device), labels] = (
            torch.max(image).detach())
        return flat.reshape_as(image)

    def ff_forward(self, image, training=False, update_stats=False):
        """Forward candidate-labelled data through normalized residual edges."""
        hiddens = []
        caches = []
        pre = self._ff_normalize(image)
        convolution = F.conv2d(pre, self.W[0], stride=1, padding=1)
        normalized, norm_cache = self._normalize(
            0, convolution, training, update_stats)
        hidden = F.relu(normalized)
        hiddens.append(hidden)
        caches.append({
            "pre": pre, "normalization": norm_cache, "shortcut": None})
        for block in self.blocks:
            first = block["first"]
            second = block["second"]
            parent = hidden
            pre = self._ff_normalize(parent)
            convolution = F.conv2d(
                pre, self.W[first], stride=block["stride"], padding=1)
            normalized, norm_cache = self._normalize(
                first, convolution, training, update_stats)
            first_hidden = F.relu(normalized)
            hiddens.append(first_hidden)
            caches.append({
                "pre": pre, "normalization": norm_cache, "shortcut": None})

            pre = self._ff_normalize(first_hidden)
            shortcut = self._option_a_shortcut(
                parent, block["out_channels"], block["stride"])
            convolution = F.conv2d(
                pre, self.W[second], stride=1, padding=1)
            normalized, norm_cache = self._normalize(
                second, convolution, training, update_stats)
            hidden = F.relu(
                shortcut + self.residual_scale * normalized)
            hiddens.append(hidden)
            caches.append({
                "pre": pre, "normalization": norm_cache,
                "shortcut": shortcut})
        features = self._ff_normalize(hidden.mean(dim=(2, 3)))
        logits = features @ self.W_out.t() + self.b_out
        return {
            "hiddens": hiddens,
            "caches": caches,
            "features": features,
            "logits": logits,
        }

    def _ff_local_outputs(self, layer_index, positive, negative):
        """Re-evaluate exactly one target layer on detached prefix states."""
        if layer_index == self.n_hidden:
            inputs = torch.cat(
                [positive["features"], negative["features"]], dim=0).detach()
            outputs = inputs @ self.W_out.t() + self.b_out
        else:
            positive_cache = positive["caches"][layer_index]
            negative_cache = negative["caches"][layer_index]
            inputs = torch.cat(
                [positive_cache["pre"], negative_cache["pre"]],
                dim=0).detach()
            spec = self.layer_specs[layer_index]
            convolution = F.conv2d(
                inputs, self.W[layer_index],
                stride=spec.stride, padding=spec.padding)
            normalized, _ = self._normalize(
                layer_index, convolution, training=True, update_stats=True)
            if positive_cache["shortcut"] is None:
                outputs = F.relu(normalized)
            else:
                shortcut = torch.cat([
                    positive_cache["shortcut"],
                    negative_cache["shortcut"],
                ], dim=0).detach()
                outputs = F.relu(
                    shortcut + spec.branch_scale * normalized)
        return outputs.chunk(2, dim=0)

    def ff_train_layer(self, layer_index, image, labels,
                       learning_rate=0.03, negative_labels=None):
        """Update one greedy FF layer; every prefix and non-target is detached."""
        if not 0 <= layer_index < self.ff_num_layers:
            raise ValueError("invalid Forward-Forward layer")
        if negative_labels is None:
            offsets = torch.randint(
                1, self.n_classes, labels.shape, device=labels.device)
            negative_labels = (labels + offsets) % self.n_classes
        with torch.no_grad():
            positive = self.ff_forward(
                self.ff_overlay(image, labels),
                training=True, update_stats=False)
            negative = self.ff_forward(
                self.ff_overlay(image, negative_labels),
                training=True, update_stats=False)
        all_parameters = (
            self.W + self.gamma + self.beta + [self.W_out, self.b_out])
        for parameter in all_parameters:
            parameter.grad = None
        optimizer = self.ff_optimizers[layer_index]
        optimizer.param_groups[0]["lr"] = learning_rate
        positive_output, negative_output = self._ff_local_outputs(
            layer_index, positive, negative)
        axes = tuple(range(1, positive_output.ndim))
        positive_goodness = positive_output.square().mean(dim=axes)
        negative_goodness = negative_output.square().mean(dim=axes)
        loss = (
            F.softplus(-positive_goodness + self.ff_threshold)
            + F.softplus(negative_goodness - self.ff_threshold)
        ).mean()
        loss.backward()
        optimizer.step()
        return {
            "loss": float(loss.detach()),
            "positive_goodness": float(positive_goodness.mean().detach()),
            "negative_goodness": float(negative_goodness.mean().detach()),
            "pair_accuracy": float(
                (positive_goodness > negative_goodness).float().mean()),
        }

    @torch.no_grad()
    def ff_candidate_scores(self, image):
        scores = []
        for candidate in range(self.n_classes):
            labels = torch.full(
                (image.shape[0],), candidate, device=image.device,
                dtype=torch.long)
            forward = self.ff_forward(
                self.ff_overlay(image, labels), training=False)
            layer_goodness = [
                value.square().mean(dim=tuple(range(1, value.ndim)))
                for value in forward["hiddens"] + [forward["logits"]]]
            scores.append(sum(
                layer_goodness[self.ff_score_from_layer:]))
        return torch.stack(scores, dim=1)


class CIFARDualPropResNet(CIFARLocalResNet):
    """Dual Propagation on the residual DAG with author DP-transpose updates.

    ``s_plus`` and ``s_minus`` are initialized to the ordinary forward states.
    A ``fwK`` pass updates every residual-DAG node in topological order.  The
    feedforward drive uses the forward edge, while the difference of each
    child state is transported through the exact transpose of that same edge.
    This is intentional symmetric feedback in the Dual Propagation baseline,
    not a claim of weight-transport-free learning.
    """

    def __init__(self, *args, alpha=0.0, dp_beta=0.1,
                 inference_passes=16, **kwargs):
        super().__init__(*args, **kwargs)
        if not 0.0 <= alpha <= 1.0:
            raise ValueError("Dual Propagation alpha must lie in [0, 1]")
        if dp_beta <= 0 or inference_passes < 1:
            raise ValueError("invalid Dual Propagation inference settings")
        self.dp_alpha = float(alpha)
        self.dp_beta = float(dp_beta)
        self.dp_inference_passes = int(inference_passes)
        self._node_kind = {0: ("stem",)}
        self._outgoing = {index: [] for index in range(self.n_hidden)}
        for block in self.blocks:
            first = block["first"]
            second = block["second"]
            parent = first - 1
            self._node_kind[first] = ("first", parent)
            self._node_kind[second] = (
                "second", first, parent, block["out_channels"],
                block["stride"])
            self._outgoing[parent].append((first, "conv"))
            self._outgoing[first].append((second, "conv"))
            self._outgoing[parent].append((second, "shortcut"))

    @staticmethod
    def _option_a_shortcut_transpose(value, input_shape, stride):
        """Adjoint of the parameter-free option-A shortcut."""
        in_channels = input_shape[1]
        out_channels = value.shape[1]
        missing = out_channels - in_channels
        if missing < 0:
            raise ValueError("option-A transpose cannot increase input width")
        before = missing // 2
        selected = value[:, before:before + in_channels]
        if stride == 1:
            if tuple(selected.shape) != tuple(input_shape):
                raise ValueError("option-A transpose shape mismatch")
            return selected
        result = value.new_zeros(input_shape)
        result[:, :, ::2, ::2] = selected
        return result

    def _node_prediction(self, index, states, image):
        """Return the unrectified local prediction and its edge cache."""
        kind = self._node_kind[index]
        if kind[0] == "stem":
            pre = image
            shortcut = None
        elif kind[0] == "first":
            pre = states[kind[1]]
            shortcut = None
        else:
            pre = states[kind[1]]
            shortcut = self._option_a_shortcut(
                states[kind[2]], kind[3], kind[4])
        spec = self.layer_specs[index]
        convolution = F.conv2d(
            pre, self.W[index], stride=spec.stride, padding=spec.padding)
        normalized, normalization = self._normalize(
            index, convolution, training=True, update_stats=False)
        prediction = (
            shortcut + spec.branch_scale * normalized
            if shortcut is not None else normalized)
        return prediction, {
            "pre": pre,
            "normalization": normalization,
            "shortcut": shortcut,
        }

    def _edge_transpose(self, child, edge_kind, field, states, image):
        """Apply one residual-DAG edge transpose to a child state field."""
        kind = self._node_kind[child]
        if edge_kind == "shortcut":
            if kind[0] != "second":
                raise AssertionError("only a second convolution has a shortcut")
            return self._option_a_shortcut_transpose(
                field, states[kind[2]].shape, kind[4])
        _, cache = self._node_prediction(child, states, image)
        spec = self.layer_specs[child]
        local_field = field * spec.branch_scale
        local_field, _, _ = self._normalization_backward(
            child, local_field, cache["normalization"])
        return torch.nn.grad.conv2d_input(
            cache["pre"].shape, self.W[child], local_field,
            stride=spec.stride, padding=spec.padding)

    def _outgoing_feedback(self, index, deltas, states, image):
        result = torch.zeros_like(states[index])
        for child, edge_kind in self._outgoing[index]:
            result.add_(self._edge_transpose(
                child, edge_kind, deltas[child], states, image))
        if index == self.n_hidden - 1:
            spatial = states[index].shape[2] * states[index].shape[3]
            result.add_(
                (deltas[-1] @ self.W_out)[:, :, None, None] / spatial)
        return result

    def infer_dual_states(self, image, one_hot, clean_forward=None):
        """Run the author ``fwK`` DP-transpose state updates on the DAG."""
        if clean_forward is None:
            clean_forward = self.forward(
                image, return_cache=True, training=True, update_stats=False)
        plus = [
            value.detach().clone() for value in clean_forward["hiddens"]]
        minus = [value.detach().clone() for value in plus]
        plus.append(clean_forward["logits"].detach().clone())
        minus.append(clean_forward["logits"].detach().clone())
        fixed_prediction = clean_forward["logits"].detach()
        alpha = self.dp_alpha
        for _ in range(self.dp_inference_passes):
            for index in range(self.n_hidden):
                states = [
                    alpha * positive + (1.0 - alpha) * negative
                    for positive, negative in zip(plus[:-1], minus[:-1])]
                prediction, _ = self._node_prediction(index, states, image)
                deltas = [
                    positive - negative
                    for positive, negative in zip(plus, minus)]
                feedback = self._outgoing_feedback(
                    index, deltas, states, image)
                plus[index] = F.relu(
                    prediction + (1.0 - alpha) * feedback)
                minus[index] = F.relu(prediction - alpha * feedback)
            states = [
                alpha * positive + (1.0 - alpha) * negative
                for positive, negative in zip(plus[:-1], minus[:-1])]
            features = states[-1].mean(dim=(2, 3))
            prediction = features @ self.W_out.t() + self.b_out
            output_field = self.dp_beta * (
                torch.softmax(fixed_prediction, dim=1) - one_hot)
            plus[-1] = prediction - (1.0 - alpha) * output_field
            minus[-1] = prediction + alpha * output_field
        return plus, minus

    @torch.no_grad()
    def dualprop_ascent_directions(self, image, plus, minus):
        """Evaluate the local DP contrastive correlations without autograd."""
        alpha = self.dp_alpha
        beta = self.dp_beta
        states = [
            alpha * positive + (1.0 - alpha) * negative
            for positive, negative in zip(plus[:-1], minus[:-1])]
        deltas = [
            (positive - negative) / beta
            for positive, negative in zip(plus, minus)]
        directions = []
        gamma_directions = []
        beta_directions = []
        batch = image.shape[0]
        for index in range(self.n_hidden):
            _, cache = self._node_prediction(index, states, image)
            spec = self.layer_specs[index]
            local_field = deltas[index] * spec.branch_scale
            local_field, gamma_direction, beta_direction = (
                self._normalization_backward(
                    index, local_field, cache["normalization"]))
            direction = torch.nn.grad.conv2d_weight(
                cache["pre"], self.W[index].shape, local_field,
                stride=spec.stride, padding=spec.padding)
            directions.append(direction / batch)
            if gamma_direction is not None:
                gamma_directions.append(gamma_direction / batch)
                beta_directions.append(beta_direction / batch)
        features = states[-1].mean(dim=(2, 3))
        output_weight = deltas[-1].t() @ features / batch
        output_bias = deltas[-1].mean(dim=0)
        return (
            directions, gamma_directions, beta_directions,
            output_weight, output_bias)

    def dualprop_step(self, image, labels, eta, eta_output=None,
                      momentum=0.0, weight_decay=0.0):
        """One fully local DP-transpose update on a matched ResNet batch."""
        one_hot = F.one_hot(labels, self.n_classes).to(image.dtype)
        with torch.no_grad():
            clean = self.forward(
                image, return_cache=True, training=True, update_stats=True)
            loss = F.cross_entropy(clean["logits"], labels)
            plus, minus = self.infer_dual_states(
                image, one_hot, clean_forward=clean)
            (directions, gamma_directions, beta_directions,
             output_weight, output_bias) = self.dualprop_ascent_directions(
                 image, plus, minus)
            self.apply_ascent(
                directions, output_weight, output_bias, eta,
                eta_output=eta_output, momentum=momentum,
                weight_decay=weight_decay,
                gamma_directions=gamma_directions,
                beta_directions=beta_directions)
        return float(loss)


class CIFAREquilibriumPropResNet(CIFARDualPropResNet):
    """Canonical two-phase EP dynamics on the matched residual energy graph.

    The Dual-Propagation parent class is used only for its audited residual-edge
    operators and exact edge transposes.  EP itself has one hard-sigmoid state
    per forward population, synchronous leaky relaxation, random nudge sign,
    and free-versus-nudged local correlations.
    """

    def __init__(self, *args, ep_beta=0.5, dt=0.5, free_steps=20,
                 nudge_steps=4, random_beta_sign=True, **kwargs):
        super().__init__(
            *args, alpha=0.0, dp_beta=0.1, inference_passes=1, **kwargs)
        if ep_beta <= 0 or not 0 < dt <= 1:
            raise ValueError("invalid Equilibrium Propagation scale or step")
        if free_steps < 1 or nudge_steps < 1:
            raise ValueError("EP phases require positive relaxation steps")
        self.ep_beta = float(ep_beta)
        self.ep_dt = float(dt)
        self.ep_free_steps = int(free_steps)
        self.ep_nudge_steps = int(nudge_steps)
        self.ep_random_beta_sign = bool(random_beta_sign)

    @staticmethod
    def ep_rho(state):
        return state.clamp(0.0, 1.0)

    @staticmethod
    def ep_rhop(state):
        return ((state >= 0.0) & (state <= 1.0)).to(state.dtype)

    @torch.no_grad()
    def ep_settle(self, image, one_hot, beta=0.0, steps=None,
                  initial_states=None):
        steps = self.ep_free_steps if steps is None else int(steps)
        clipped_image = self.ep_rho(image)
        if initial_states is None:
            shapes = self.forward(
                clipped_image, training=True,
                update_stats=False)["hiddens"]
            states = [torch.zeros_like(value) for value in shapes]
            states.append(torch.zeros(
                image.shape[0], self.n_classes,
                device=image.device, dtype=image.dtype))
        else:
            states = [value.detach().clone() for value in initial_states]
        for _ in range(steps):
            rho_hidden = [
                self.ep_rho(value) for value in states[:-1]]
            fields = rho_hidden + [self.ep_rho(states[-1])]
            updated = []
            for index, state in enumerate(states[:-1]):
                prediction, _ = self._node_prediction(
                    index, rho_hidden, clipped_image)
                feedback = self._outgoing_feedback(
                    index, fields, rho_hidden, clipped_image)
                drive = -self.ep_rho(state) + prediction + feedback
                updated.append(self.ep_rho(
                    state + self.ep_dt * self.ep_rhop(state) * drive))
            features = rho_hidden[-1].mean(dim=(2, 3))
            prediction = features @ self.W_out.t() + self.b_out
            output = states[-1]
            drive = -self.ep_rho(output) + prediction
            if beta:
                drive = drive + 2.0 * beta * (one_hot - output)
            updated.append(self.ep_rho(
                output + self.ep_dt * self.ep_rhop(output) * drive))
            states = updated
        return states

    @torch.no_grad()
    def _ep_phase_correlations(self, image, states):
        clipped_image = self.ep_rho(image)
        rho_hidden = [self.ep_rho(value) for value in states[:-1]]
        batch = image.shape[0]
        directions = []
        gamma_directions = []
        beta_directions = []
        for index in range(self.n_hidden):
            _, cache = self._node_prediction(
                index, rho_hidden, clipped_image)
            spec = self.layer_specs[index]
            field = self.ep_rho(states[index]) * spec.branch_scale
            field, gamma_direction, beta_direction = (
                self._normalization_backward(
                    index, field, cache["normalization"]))
            direction = torch.nn.grad.conv2d_weight(
                cache["pre"], self.W[index].shape, field,
                stride=spec.stride, padding=spec.padding)
            directions.append(direction / batch)
            if gamma_direction is not None:
                gamma_directions.append(gamma_direction / batch)
                beta_directions.append(beta_direction / batch)
        features = rho_hidden[-1].mean(dim=(2, 3))
        output = self.ep_rho(states[-1])
        output_weight = output.t() @ features / batch
        output_bias = output.mean(dim=0)
        return (
            directions, gamma_directions, beta_directions,
            output_weight, output_bias)

    @torch.no_grad()
    def ep_ascent_directions(self, image, free_states, nudged_states,
                             signed_beta):
        if signed_beta == 0:
            raise ValueError("EP contrast requires a nonzero nudge")
        free = self._ep_phase_correlations(image, free_states)
        nudged = self._ep_phase_correlations(image, nudged_states)
        result = []
        for nudged_group, free_group in zip(nudged, free):
            if isinstance(nudged_group, list):
                result.append([
                    (nudged_value - free_value) / signed_beta
                    for nudged_value, free_value in zip(
                        nudged_group, free_group)])
            else:
                result.append((nudged_group - free_group) / signed_beta)
        return tuple(result)

    def ep_step(self, image, labels, eta, eta_output=None,
                momentum=0.0, weight_decay=0.0, generator=None):
        one_hot = F.one_hot(labels, self.n_classes).to(image.dtype)
        with torch.no_grad():
            free_states = self.ep_settle(
                image, one_hot, beta=0.0, steps=self.ep_free_steps)
            sign = 1.0
            if self.ep_random_beta_sign:
                draw = torch.randint(
                    0, 2, (), device=image.device, generator=generator)
                sign = 1.0 if int(draw) else -1.0
            signed_beta = sign * self.ep_beta
            nudged_states = self.ep_settle(
                image, one_hot, beta=signed_beta,
                steps=self.ep_nudge_steps, initial_states=free_states)
            (directions, gamma_directions, beta_directions,
             output_weight, output_bias) = self.ep_ascent_directions(
                 image, free_states, nudged_states, signed_beta)
            self.apply_ascent(
                directions, output_weight, output_bias, eta,
                eta_output=eta_output, momentum=momentum,
                weight_decay=weight_decay,
                gamma_directions=gamma_directions,
                beta_directions=beta_directions)
            loss = F.mse_loss(free_states[-1], one_hot)
        return float(loss), free_states