pyjuice.compile

pyjuice.compile(ns: CircuitNodes, layer_sparsity_tol: float = 0.5, max_num_partitions: int | None = None, disable_gpu_compilation: bool = False, force_gpu_compilation: bool = False, max_tied_ns_per_parflow_block: int = 8, verbose: bool = True) → Module

Compile a PC represented by a DAG into an equivalent torch.nn.Module.

Parameters:

• ns (CircuitNodes) – the root node of the PC’s DAG

• layer_sparsity_tol (float) – the maximum allowed fraction for added pseudo edges within every layer (better to set to a small number for sparse/block-sparse PCs)

• max_num_partitions (Optional[int]) – maximum number of partitions in a layer

• disable_gpu_compilation (bool) – force PyJuice to use CPU compilation

• force_gpu_compilation (bool) – force PyJuice to use GPU compilation

• max_tied_ns_per_parflow_block (int) – how many groups of tied parameters are allowed to share the same flow/gradient accumulator (higher values -> consumes less GPU memory; lower values -> potentially avoid stalls caused by atomic operations)

• verbose (bool) – Whether to display the progress of the compilation

Returns:

the compiled PC with type torch.nn.Module