nd2py.utils.nn package

nd2py.utils.nn.setup_lazy_imports(module_name: str, import_mapping: Dict[str, Tuple[str, str]])

Set up lazy imports for a module’s __init__.py.

Returns (__getattr__, __dir__, __all__) which should be assigned at the module level so that from package import OptionalClass works without importing the optional dependency until it is actually needed.

Parameters:

• module_name – The __name__ of the calling module.

• import_mapping – A dict mapping attribute names to (module_path, requires) tuples. module_path is a relative import path (e.g. ".torch_calc") and requires is the optional-dependency group name (e.g. "nn") shown in the error message when the dependency is missing.

Usage:

# __init__.py
from .core import CoreClass
from ..utils.lazy_loader import setup_lazy_imports

if TYPE_CHECKING:
    from .optional import OptionalClass

__getattr__, __dir__, __all__ = setup_lazy_imports(__name__, {
    "OptionalClass": (".optional", "nn"),
})

Submodules

nd2py.utils.nn.gnn module

class nd2py.utils.nn.gnn.GNNLayer(*args: Any, **kwargs: Any)

Bases: Module

__init__(node_dim, edge_dim, out_dim, residual=False)

forward(v, e, edge_list, num_nodes)

• v: [N, V, node_dim] (N: SampleNum, V: TotalNodes)

• e: [N, E, edge_dim] (N: SampleNum, E: TotalEdges)

• edge_list: [2, E] (第0行: source, 第1行: target)

• num_nodes: [N] (每个样本的节点数量)

• return: [N, V, out_dim], [N, E, out_dim]

class nd2py.utils.nn.gnn.GNN(*args: Any, **kwargs: Any)

Bases: Module

__init__(d_emb, n_layers, node_dim, edge_dim, dropout)

fwd(v, e, edge_list, num_nodes)

forward(v, e, edge_list, num_nodes, chunk_size=None)

nd2py.utils.nn.positional_encoding module

class nd2py.utils.nn.positional_encoding.PositionalEncoding(*args: Any, **kwargs: Any)

Bases: Module

__init__(d_emb, dropout=0.2, max_len=5000)

forward(x)