spb/wpai0001_uqo

1"""Gate C for the MLX backend: wrapped layers are a bitwise no-op when idle; interventions change logits.2Skipped when MLX or the checkpoint is unavailable (never fakes support)."""3import os45import numpy as np6import pytest78mlx = pytest.importorskip("mlx.core")9pytest.importorskip("mlx_lm")1011MODEL = os.environ.get("WP_AI_0001_MLX_TEST_MODEL", "mlx-community/Qwen2.5-1.5B-Instruct-bf16")121314@pytest.mark.skipif(os.environ.get("WP_AI_0001_SKIP_MLX") == "1", reason="MLX test disabled")15def test_wrapped_layers_are_bitwise_noop_and_steering_changes_output():16    from huggingface_hub import try_to_load_from_cache17    from mlx_lm import load1819    from wp_ai_0001.models.mlx_backend import MLXAddAtLast, MLXCapture, instrument2021    if try_to_load_from_cache(MODEL, "config.json") is None:22        pytest.skip("checkpoint not cached")23    model, tok = load(MODEL)24    ids = mlx.array([tok.encode("The capital of France is")])25    ref = np.array(model(ids)[0, -1, :].astype(mlx.float32))26    hooks = instrument(model)27    out = np.array(model(ids)[0, -1, :].astype(mlx.float32))28    assert np.array_equal(ref, out)29    cap = MLXCapture("last")30    hooks.add("resid_post.10", cap)31    out2 = np.array(model(ids)[0, -1, :].astype(mlx.float32))32    assert np.array_equal(ref, out2) and "resid_post.10" in cap.store33    hooks.clear()34    hooks.add("resid_post.10", MLXAddAtLast(np.ones(model.args.hidden_size, dtype=np.float32) * 5.0))35    hooks.enabled = False36    out3 = np.array(model(ids)[0, -1, :].astype(mlx.float32))37    assert np.array_equal(ref, out3)38    hooks.enabled = True39    out4 = np.array(model(ids)[0, -1, :].astype(mlx.float32))40    assert not np.array_equal(ref, out4)41