torchcell.transforms.hetero_to_dense_mask.HeteroToDenseMask

class HeteroToDenseMask(num_nodes_dict: dict[str, int] | None = None, square_attrs: dict[str, list[str]] | None = None)

Bases: BaseTransform

Convert sparse hetero adjacencies to boolean dense masks.

Edge attributes are preserved in their original sparse format for memory efficiency while adjacency/incidence matrices become boolean dense masks.

A mask entry is set only when both endpoints are real (original) nodes, so an edge into a padding row or past the input’s node count is dropped from the mask (it stays in edge_index). A node tensor whose first axis equals the original node count N is padded with zero rows. Only the attributes named in square_attrs are also padded on their second axis, so a listed per-node [N, N] matrix becomes [N_pad, N_pad]; an unlisted [N, d] tensor keeps its d columns even when d == N (the shape cannot tell a feature axis from a node axis).

Parameters:

• num_nodes_dict (Dict[str, int], optional) – Dictionary mapping node types to their desired number of nodes. If not provided for a node type, will use the maximum number of nodes found for that type. (default: None)

• square_attrs (Dict[str, List[str]], optional) – Per node type, the names of the [N, N, ...] node attributes to pad on both axes. A listed attribute that is absent or not [N, N, ...], or a node type the data does not have, raises ValueError in forward; a reserved name (RESERVED_NODE_ATTRS or a leading underscore) raises ValueError at construction. (default: None, no square attributes)

forward(data: HeteroData) → HeteroData

Add dense masks, pad node attributes, and return the transformed data.