Skip to content

Input and Adapter Types

types

BatchParam dataclass

Bases: Generic[_BatchT]

Explicit shared/separate wrapper for batch inputs.

Source code in src/mifrost/encoders/types.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
@dataclass(frozen=True)
class BatchParam(Generic[_BatchT]):
    """Explicit shared/separate wrapper for batch inputs."""

    kind: Literal["shared", "separate", "none"]
    value: Any | None = None

    @classmethod
    def shared(cls, value: _BatchT) -> "BatchParam[_BatchT]":
        return cls(kind="shared", value=value)

    @classmethod
    def separate(cls, values: Iterable[_BatchT | None]) -> "BatchParam[_BatchT]":
        return cls(kind="separate", value=list(values))

    @classmethod
    def none(cls) -> "BatchParam[_BatchT]":
        return cls(kind="none", value=None)

BatchEncodingLike

Bases: Protocol

Structural type for native C++ BatchEncoding-like objects.

Concrete bindings expose this shape from C++ while Python code can type against the protocol without importing binding internals. The protocol mirrors PyG-level structural metadata (graph/node/edge counts and type lists) but intentionally does not model dynamic node storage attributes like data["atom"].x.

Source code in src/mifrost/encoders/types.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@runtime_checkable
class BatchEncodingLike(Protocol):
    """
    Structural type for native C++ ``BatchEncoding``-like objects.

    Concrete bindings expose this shape from C++ while Python code can
    type against the protocol without importing binding internals.
    The protocol mirrors PyG-level structural metadata (graph/node/edge
    counts and type lists) but intentionally does not model dynamic node
    storage attributes like ``data["atom"].x``.
    """

    num_graphs: int
    num_nodes: int
    num_edges: int
    graph_kind: str
    node_types: Sequence[str]
    edge_types: Sequence[tuple[str, str, str]]

    def as_dict(self) -> EncodingDict: ...

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

    def as_hetero(self) -> "HeteroBatchEncodingView": ...

    def as_homo(self) -> "HomoBatchEncodingView": ...

    def to(self, device: Any) -> "BatchEncodingLike": ...

    def schema_fingerprint(self) -> int: ...

HeteroEncoding

Bases: BatchEncodingLike, Protocol

Refined protocol for hetero native encodings.

Source code in src/mifrost/encoders/types.py 
127
128
129
130
131
132
133
@runtime_checkable
class HeteroEncoding(BatchEncodingLike, Protocol):
    """Refined protocol for hetero native encodings."""

    graph_kind: Literal["hetero"]

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

HomoEncoding

Bases: BatchEncodingLike, Protocol

Refined protocol for homo native encodings.

Source code in src/mifrost/encoders/types.py 
136
137
138
139
140
141
142
@runtime_checkable
class HomoEncoding(BatchEncodingLike, Protocol):
    """Refined protocol for homo native encodings."""

    graph_kind: Literal["homo"]

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

FlatEncoding

Bases: BatchEncodingLike, Protocol

Refined protocol for flat native encodings.

Flat encodings use graph_kind == "flat" even though their tensor layout is still homo-shaped and converts through the flat PyG carrier.

Source code in src/mifrost/encoders/types.py 
145
146
147
148
149
150
151
152
153
154
155
156
@runtime_checkable
class FlatEncoding(BatchEncodingLike, Protocol):
    """
    Refined protocol for flat native encodings.

    Flat encodings use `graph_kind == "flat"` even though their tensor layout is
    still homo-shaped and converts through the flat PyG carrier.
    """

    graph_kind: Literal["flat"]

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

register_state_adapter(state_type, adapter)

Register a legacy exact-type Pymimir state adapter.

Source code in src/mifrost/encoders/types.py 
168
169
170
def register_state_adapter(state_type: type[object], adapter: Any) -> None:
    """Register a legacy exact-type Pymimir state adapter."""
    _pymimir_types().register_state_adapter(state_type, adapter)

register_domain_adapter(domain_type, adapter)

Register a legacy exact-type Pymimir domain adapter.

Source code in src/mifrost/encoders/types.py 
185
186
187
def register_domain_adapter(domain_type: type[object], adapter: Any) -> None:
    """Register a legacy exact-type Pymimir domain adapter."""
    _pymimir_types().register_domain_adapter(domain_type, adapter)

register_literal_adapter(literal_type, adapter)

Register a legacy exact-type Pymimir literal adapter.

Source code in src/mifrost/encoders/types.py 
202
203
204
def register_literal_adapter(literal_type: type[object], adapter: Any) -> None:
    """Register a legacy exact-type Pymimir literal adapter."""
    _pymimir_types().register_literal_adapter(literal_type, adapter)

register_action_adapter(action_type, adapter)

Register a legacy exact-type Pymimir action adapter.

Source code in src/mifrost/encoders/types.py 
219
220
221
def register_action_adapter(action_type: type[object], adapter: Any) -> None:
    """Register a legacy exact-type Pymimir action adapter."""
    _pymimir_types().register_action_adapter(action_type, adapter)