Skip to content

compute_gradiend_transition_cross_encoding_matrix

compute_gradiend_transition_cross_encoding_matrix

compute_gradiend_transition_cross_encoding_matrix(trainers, *, trainer_order=None, transition_order=None, encoder_summary=None, split='test', max_size=None, use_cache=None, seed_selection=None, seed_aggregate='mean', dispersion=None)

Compute a GRADIEND × directed-transition cross-encoding matrix.

Each row is one trained GRADIEND; each column is a directed input transition factual_class->alternative_class from the shared cross-task test pool. Cell (i, j) is the mean encoded response of GRADIEND i on snippets with transition j (within-model average only).

This is the standard rectangular matrix before anchor sign alignment and anchor-class aggregation used by :func:compute_anchor_aligned_encoding_matrix.

Source code in gradiend/comparison/cross_encoding.py
def compute_gradiend_transition_cross_encoding_matrix(
    trainers: Dict[str, object],
    *,
    trainer_order: Optional[Sequence[str]] = None,
    transition_order: Optional[Sequence[str]] = None,
    encoder_summary: Optional[Dict[str, Any]] = None,
    split: str = "test",
    max_size: Optional[int] = None,
    use_cache: Optional[bool] = None,
    seed_selection: Optional[str] = None,
    seed_aggregate: str = "mean",
    dispersion: Optional[str] = None,
) -> Dict[str, Any]:
    """Compute a GRADIEND × directed-transition cross-encoding matrix.

    Each row is one trained GRADIEND; each column is a directed input transition
    ``factual_class->alternative_class`` from the shared cross-task test pool.
    Cell ``(i, j)`` is the mean encoded response of GRADIEND *i* on snippets with
    transition *j* (within-model average only).

    This is the standard rectangular matrix **before** anchor sign alignment and
    anchor-class aggregation used by :func:`compute_anchor_aligned_encoding_matrix`.
    """
    if not trainers:
        raise ValueError("trainers must be a non-empty dict")
    order = [str(t) for t in (trainer_order or list(trainers.keys())) if str(t) in trainers]
    if not order:
        raise ValueError("trainer_order produced no valid trainer ids")

    if encoder_summary is None:
        encoder_summary = build_cross_task_encoder_summary(
            trainers,
            [],
            split=split,
            max_size=max_size,
            use_cache=use_cache,
            seed_selection=seed_selection,
            seed_aggregate=seed_aggregate,
            dispersion=dispersion,
        )

    meta = comparison_seed_metadata(
        trainers,
        seed_selection=seed_selection,
        seed_aggregate=seed_aggregate,
        dispersion=dispersion,
    )

    per_trainer_means: Dict[str, Dict[str, Tuple[float, int]]] = {}
    per_trainer_stds: Dict[str, Dict[str, float]] = {}
    for trainer_id, payload in encoder_summary.items():
        if str(trainer_id) not in order:
            continue
        entry = payload if isinstance(payload, dict) else {}
        means, stds = _transition_stats_from_encoder_summary(entry)
        per_trainer_means[str(trainer_id)] = means
        per_trainer_stds[str(trainer_id)] = stds

    if transition_order is None:
        observed: set[str] = set()
        for values in per_trainer_means.values():
            observed.update(values.keys())
        columns = sorted(observed)
    else:
        columns = [
            normalize_transition_id(value) or str(value) for value in transition_order
        ]

    matrix: List[List[float]] = []
    n_matrix: List[List[int]] = []
    cell_stats: List[List[Dict[str, Any]]] = []
    has_std = any(stds for stds in per_trainer_stds.values())
    for trainer_id in order:
        by_transition = per_trainer_means.get(trainer_id, {})
        by_std = per_trainer_stds.get(trainer_id, {})
        row_values: List[float] = []
        row_counts: List[int] = []
        stats_row: List[Dict[str, Any]] = []
        for transition_id in columns:
            cell = by_transition.get(transition_id)
            if cell is None:
                row_values.append(float("nan"))
                row_counts.append(0)
                stats_row.append({})
            else:
                row_values.append(float(cell[0]))
                row_counts.append(int(cell[1]))
                stat: Dict[str, Any] = {"aggregate": float(cell[0]), "n": int(cell[1])}
                std_value = by_std.get(transition_id)
                if std_value is not None:
                    stat["std"] = float(std_value)
                stats_row.append(stat)
        matrix.append(row_values)
        n_matrix.append(row_counts)
        if has_std:
            cell_stats.append(stats_row)

    payload: Dict[str, Any] = {
        "measure": "gradiend_transition_cross_encoding_mean",
        "model_ids": order,
        "column_ids": columns,
        "rows": order,
        "columns": columns,
        "matrix": matrix,
        "n_matrix": n_matrix,
        "split": split,
        "max_size": max_size,
        "transitions_found": sorted(
            {t for values in per_trainer_means.values() for t in values}
        ),
        **meta,
    }
    if has_std:
        payload["cell_stats"] = cell_stats
        payload["multi_seed"] = True
    return payload