torchcell.models.dcell_opt.DCellOpt

class DCellOpt(hetero_data: HeteroData, min_subsystem_size: int = 20, subsystem_ratio: float = 0.3, output_size: int = 1)[source]

Bases: Module

Optimized DCell with tensorized operations for torch.compile.

enable_profiling(enable: bool = True) → None[source]

Enable or disable performance profiling during forward pass.

property num_parameters: dict[str, int]

Count parameters in different parts of the model.

forward(cell_graph: HeteroData, batch: HeteroData) → tuple[Tensor, dict[str, Any]][source]

Forward pass through DCell hierarchy with parallel stratum processing.