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 = 32, device: int | device | None = None, verbose: bool = True, cuda_graphs: bool = False, precision: str = 'auto') → 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)

  • device (Optional[Union[int,torch.device]]) – Which GPU do we use for compilation (the default is torch.cuda.current_device)

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

  • cuda_graphs (bool) – record CUDA graphs of the inner layers by default (see TensorCircuit); OFF by default

  • precision (str) – precision of the sum layers’ products: “auto” (the default; the fastest per layer), “tf32” or “fp32” (see TensorCircuit.precision; settable after compiling)

Returns:

the compiled PC with type torch.nn.Module