Listing: Listing

def repeat_point(x, length):
    if not isinstance(length, int) or length < 1:
        raise ValueError("positive sequence length required")
    return x[:, None, :].expand(-1, length, -1)

def normal_cdf(x):
    return 0.5 * (1.0 + math.erf(x / math.sqrt(2.0)))

def fresh_paired_probes(n, length, seed, dtype=torch.float64):
    if not isinstance(n, int) or n < 1:
        raise ValueError("positive example count required")
    gen = generator(seed, "fresh_probe_latents")
    centers = torch.tensor(
        [[-3., -3.], [3., 3.], [-3., 3.], [3., -3.]],
        dtype=dtype
    )
    y = torch.randint(0, 4, (n,), generator=gen)
    distractor_y = torch.randint(0, 4, (n,), generator=gen)
    clean, distractor = centers[y], centers[distractor_y]
    z = torch.randn((n, 2), dtype=dtype, generator=gen)
    angle = math.pi / 4
    rotation = torch.tensor(
        [[math.cos(angle), -math.sin(angle)],
         [math.sin(angle), math.cos(angle)]], dtype=dtype
    )
    raw = {
        "clean": clean,
        "standard": clean + 0.8 * z,
        "rotation_row_minus_45": clean @ rotation,
        "scale_1.5": 1.5 * clean,
        "fine_0.5_noise_0.8": 0.5 * clean + 0.8 * z,
        "input_noise_2.0": clean + 2.0 * z,
    }
    ceilings = {
        "clean": 1.0,
        "standard": normal_cdf(3.0 / 0.8)**2,
        "rotation_row_minus_45": 1.0,
        "scale_1.5": 1.0,
        "fine_0.5_noise_0.8": normal_cdf(1.5 / 0.8)**2,
        "input_noise_2.0": normal_cdf(3.0 / 2.0)**2,
    }
    for alpha in (0.25, 0.5, 0.75):
        key = f"fresh_center_blend_{alpha}"
        raw[key] = (1.0 - alpha) * clean + alpha * distractor
        ceilings[key] = 9.0 / 16.0 if alpha == 0.5 else 1.0

    return {
        key: {
            "sequence": repeat_point(value, length),
            "labels": y,
            "task_ceiling_before_extra_corruption": ceilings[key],
            "mode": "ordinary",
            "xi": 0.0, "tau": 0.0, "sigma": 0.0, "beta": 0.0,
            "design": "fresh_iid_latent_rows_shared_across_probes",
        }
        for key, value in raw.items()
    }

def delayed_cue(n, length, seed, distractor_std=0.5,
                dtype=torch.float64):
    if not isinstance(n, int) or n < 1 or length < 2:
        raise ValueError("nonempty task and at least two steps required")
    distractor_std = finite(distractor_std, "distractor_std")
    if distractor_std < 0:
        raise ValueError("negative distractor intensity")
    gen = generator(seed, "delayed_cue")
    centers = torch.tensor(
        [[-3., -3.], [3., 3.], [-3., 3.], [3., -3.]],
        dtype=dtype
    )
    labels = torch.randint(0, 4, (n,), generator=gen)
    sequence = distractor_std * torch.randn(
        (n, length, 2), dtype=dtype, generator=gen
    )
    sequence[:, 0, :] = centers[labels]
    # Subsequent inputs contain no repeated label-bearing cue.
    return sequence, labels

def battery_assessment(checkpoints, probes, seed, delta=0.05):
    delta = finite(delta, "delta")
    if not checkpoints or not probes or not 0 < delta < 1:
        raise ValueError("invalid battery specification")
    number_of_means = len(checkpoints) * len(probes)
    result = {}
    for checkpoint_id, model in checkpoints.items():
        result[checkpoint_id] = {}
        for probe_id, probe in probes.items():
            # Same per-probe seed couples draws across checkpoints.
            probe_seed_gen = generator(seed, "probe:" + probe_id)
            probe_seed = int(torch.randint(
                0, 2**31 - 1, (1,), generator=probe_seed_gen
            ))
            row = evaluate_explicit(
                model, probe["sequence"], probe["labels"], probe_seed,
                mode=probe["mode"], xi=probe["xi"], tau=probe["tau"],
                sigma=probe["sigma"], beta=probe["beta"]
            )
            n = row["n_examples"]
            radius = math.sqrt(
                math.log(2.0 * number_of_means / delta) / (2.0 * n)
            )
            row["simultaneous_radius"] = radius
            row["interval"] = (
                max(0.0, row["accuracy"] - radius),
                min(1.0, row["accuracy"] + radius)
            )
            row["checkpoint_stage"] = str(checkpoint_id)
            row["probe_design"] = probe["design"]
            row["task_ceiling_before_extra_corruption"] = \
                probe.get("task_ceiling_before_extra_corruption")
            result[checkpoint_id][probe_id] = row
    return result