Function mlp_forward

tests/dump_phase6_reference.py:233–246

(x: torch.Tensor,
                weights: Dict[str, torch.Tensor],
                prefix: str,
                num_layers: int,
                sigmoid_output: bool = False)

def mlp_forward(x: torch.Tensor,
                weights: Dict[str, torch.Tensor],
                prefix: str,
                num_layers: int,
                sigmoid_output: bool = False) -> torch.Tensor:
    for i in range(num_layers):
        w = weights[f"{prefix}.layers.{i}.weight"].float()
        b = weights[f"{prefix}.layers.{i}.bias"].float()
        x = F.linear(x, w, b)
        if i < num_layers - 1:
            x = F.relu(x)
    if sigmoid_output:
        x = torch.sigmoid(x)
    return x


def layer_norm_2d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: