Function mlp_forward

tests/dump_phase7_reference.py:261–268

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

def mlp_forward(x: torch.Tensor, weights: Dict[str, torch.Tensor], prefix: str, num_layers: int) -> 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)
    return x


def build_no_point_prompt_encoder_outputs(prompt_weights: Dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: