pyjuice.TensorCircuit
- class pyjuice.TensorCircuit(root_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)
A class for compiled PCs. It is a subclass of torch.nn.Module.
- Parameters:
root_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
- forward(inputs: Tensor, input_layer_fn: str | Callable | None = None, cache: dict | None = None, return_cache: bool = False, record_cudagraph: bool = False, apply_cudagraph: bool = True, force_use_bf16: bool = False, force_use_fp32: bool = False, propagation_alg: str | Sequence[str] | None = None, _inner_layers_only: bool = False, _no_buffer_reset: bool = False, **kwargs)
Forward evaluation of the PC.
- Parameters:
inputs (torch.Tensor) – input tensor of size [B, num_vars]
input_layer_fn (Optional[Union[str,Callable]]) – Custom forward function for input layers; if it is a string, then try to call the corresponding member function of the input layers
- backward(inputs: Tensor | None = None, ll_weights: Tensor | None = None, compute_param_flows: bool = True, flows_memory: float = 1.0, input_layer_fn: str | Callable | None = None, cache: dict | None = None, return_cache: bool = False, record_cudagraph: bool = False, apply_cudagraph: bool = True, allow_modify_flows: bool = True, propagation_alg: str | Sequence[str] = 'LL', logspace_flows: bool = False, negate_pflows: bool = False, _inner_layers_only: bool = False, _disable_buffer_init: bool = False, **kwargs)
Backward evaluation of the PC that computes node flows as well as parameter flows.
- Parameters:
inputs (torch.Tensor) – input tensor of size [B, num_vars]
ll_weights (torch.Tensor) – weights of the log-likelihoods of size [B] or [num_roots, B]
input_layer_fn (Optional[Union[str,Callable]]) – Custom forward function for input layers; if it is a string, then try to call the corresponding member function of the input layers
- mini_batch_em(step_size: float, pseudocount: float = 0.0, keep_zero_params: bool = False, step_size_rescaling: bool = False)
Perform an EM parameter update step using the accumulated parameter flows.
- Parameters:
step_size (float) – Step size - updated_params <- (1-step_size) * params + step_size * new_params
pseudocount (float) – a pseudo count added to the parameter flows
keep_zero_params (bool) – if set to True, do not add pseudocounts to zero parameters
step_size_rescaling (bool) – whether to rescale the step size by flows
- init_param_flows(flows_memory: float = 1.0, batch_size: int | None = None)
Initialize parameter flows.
- Parameters:
flows_memory (float) – the number that the current parameter flows (if any) will be multiplied by; equivalent to zeroling the flows if set to 0