compute_gradiend_feature_cross_encoding_matrix

compute_gradiend_feature_cross_encoding_matrix

compute_gradiend_feature_cross_encoding_matrix(trainers, feature_classes, *, trainer_order=None, eval_by_class=None, split='test', max_size=None)

Compute a dense GRADIEND by feature-class cross-encoding matrix.

Row i is a trained GRADIEND; column j is a feature class. Cell (i, j) is the mean encoded value when GRADIEND i encodes shared eval snippets for class j. When eval_by_class is omitted, snippets are collected from the trainers' unified data for the requested split.

Parameters:

Name	Type	Description	Default
trainers	Dict[str, object]	Mapping from trainer id to trainer object. Trainers must be able to load a model and create a gradient training dataset for the generated probe pairs.	required
feature_classes	Sequence[str]	Column order for feature classes to evaluate.	required
trainer_order	Optional[Sequence[str]]	Optional row order. Unknown ids are ignored; at least one valid id must remain.	None
eval_by_class	Optional[Dict[str, DataFrame]]	Optional precomputed mapping from feature class to a unified-data DataFrame. If omitted, it is built from trainers via :func:collect_unified_test_rows_by_feature_class.	None
split	str	Split used when collecting unified eval rows.	'test'
max_size	Optional[int]	Optional maximum examples per feature class. If set, rows are sampled with a fixed random seed before encoding.	None

Returns:

Type	Description
Dict[str, Any]	A payload with measure, model_ids, column_ids, rows,
Dict[str, Any]	columns, matrix, n_matrix, split, max_size, and
Dict[str, Any]	eval_classes_found. Missing classes produce NaN cells and count
Dict[str, Any]	0.

Raises:

Type	Description
ValueError	If no trainers/classes are provided, trainer_order has no valid trainer ids, trainer unified data is malformed, or a trainer is not binary where probe-pair construction requires it.

Source code in gradiend/comparison/cross_encoding.py

def compute_gradiend_feature_cross_encoding_matrix(
    trainers: Dict[str, object],
    feature_classes: Sequence[str],
    *,
    trainer_order: Optional[Sequence[str]] = None,
    eval_by_class: Optional[Dict[str, pd.DataFrame]] = None,
    split: str = "test",
    max_size: Optional[int] = None,
) -> Dict[str, Any]:
    """Compute a dense GRADIEND by feature-class cross-encoding matrix.

    Row *i* is a trained GRADIEND; column *j* is a feature class. Cell ``(i, j)``
    is the mean encoded value when GRADIEND *i* encodes shared eval snippets for
    class *j*. When ``eval_by_class`` is omitted, snippets are collected from the
    trainers' unified data for the requested split.

    Args:
        trainers: Mapping from trainer id to trainer object. Trainers must be
            able to load a model and create a gradient training dataset for the
            generated probe pairs.
        feature_classes: Column order for feature classes to evaluate.
        trainer_order: Optional row order. Unknown ids are ignored; at least one
            valid id must remain.
        eval_by_class: Optional precomputed mapping from feature class to a
            unified-data DataFrame. If omitted, it is built from ``trainers`` via
            :func:`collect_unified_test_rows_by_feature_class`.
        split: Split used when collecting unified eval rows.
        max_size: Optional maximum examples per feature class. If set, rows are
            sampled with a fixed random seed before encoding.

    Returns:
        A payload with ``measure``, ``model_ids``, ``column_ids``, ``rows``,
        ``columns``, ``matrix``, ``n_matrix``, ``split``, ``max_size``, and
        ``eval_classes_found``. Missing classes produce ``NaN`` cells and count
        ``0``.

    Raises:
        ValueError: If no trainers/classes are provided, ``trainer_order`` has
            no valid trainer ids, trainer unified data is malformed, or a
            trainer is not binary where probe-pair construction requires it.
    """
    if not trainers:
        raise ValueError("trainers must be a non-empty dict")
    classes = [str(c) for c in feature_classes]
    if not classes:
        raise ValueError("feature_classes must be non-empty")
    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 eval_by_class is None:
        eval_by_class = collect_unified_test_rows_by_feature_class(trainers, split=split)

    matrix: List[List[float]] = []
    n_matrix: List[List[int]] = []
    for trainer_id in order:
        trainer = trainers[trainer_id]
        model = _load_eval_model_for_trainer(trainer)
        row_values: List[float] = []
        row_counts: List[int] = []
        try:
            for class_id in classes:
                class_df = eval_by_class.get(class_id)
                if class_df is None or class_df.empty:
                    row_values.append(float("nan"))
                    row_counts.append(0)
                    continue
                sample_n = min(len(class_df), max_size) if max_size is not None else len(class_df)
                mean_val = _mean_encoded_for_feature_class(
                    trainer,
                    model,
                    class_df,
                    max_size=max_size,
                )
                row_values.append(float(mean_val) if mean_val is not None else float("nan"))
                row_counts.append(int(sample_n) if mean_val is not None else 0)
        finally:
            del model
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        matrix.append(row_values)
        n_matrix.append(row_counts)

    return {
        "measure": "gradiend_feature_cross_encoding_mean",
        "model_ids": order,
        "column_ids": classes,
        "rows": order,
        "columns": classes,
        "matrix": matrix,
        "n_matrix": n_matrix,
        "split": split,
        "max_size": max_size,
        "eval_classes_found": sorted(eval_by_class.keys()),
    }