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, device: int | device | None = None, 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
Inference
- TensorCircuit.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, pflow_temperature: float = 1.0, _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
- TensorCircuit.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, sum_layer_pre_backward_callback: Callable | None = None, sum_layer_post_backward_callback: Callable | 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, force_use_fp32: bool = False, pflow_temperature: float = 1.0, temper_eflow: 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
- TensorCircuit.set_propagation_alg(propagation_alg: str, **kwargs)
Set the default propagation algorithm used by
forward()andbackward().- Parameters:
propagation_alg (str) – the propagation algorithm; one of “LL” (log-likelihood / standard marginal inference), “MPE” (most-probable-explanation / max-product), or “GeneralLL” (an entropy-/temperature-style generalization that interpolates between “LL” and “MPE”)
For “GeneralLL”, an alpha keyword argument must be provided. The algorithm can also be selected per-call by passing propagation_alg=… directly to
forward()/backward().
Learning
- TensorCircuit.mini_batch_em(step_size: float, pseudocount: float = 0.0, keep_zero_params: bool = False, step_size_rescaling: bool = False, use_cudagraph: 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
- TensorCircuit.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
- TensorCircuit.zero_param_flows()
Zero out parameter flows.
Inspection
- TensorCircuit.get_node_mars(ns: CircuitNodes)
Retrieve the node values of ns from the previous forward pass.
- Params ns:
the target nodes
- TensorCircuit.get_node_flows(ns: CircuitNodes, **kwargs)
Retrieve the node flows of ns from the previous backward pass.
- Params ns:
the target nodes
- TensorCircuit.get_node_params(ns: CircuitNodes, clone: bool = True, **kwargs)
Retrieve the node parameters of ns.
- Params ns:
the target nodes
- Params clone:
whether to clone the parameters
- TensorCircuit.get_node_param_flows(ns: CircuitNodes, clone: bool = True, **kwargs)
Retrieve the node parameter flows of ns.
- Params ns:
the target nodes
- Params clone:
whether to clone the parameter flows
- TensorCircuit.print_statistics()
Print the statistics of the PC.
Partial Evaluation
- TensorCircuit.enable_partial_evaluation(scopes: Sequence[BitSet] | Sequence[int], forward: bool = False, backward: bool = False, overwrite: bool = False)
Restrict subsequent forward and/or backward passes to only the nodes whose scope is contained in scopes. This speeds up repeated queries that touch a fixed subset of variables (e.g., evaluating or updating only part of the circuit). Call
disable_partial_evaluation()to revert to evaluating the whole circuit.- Parameters:
scopes (Union[Sequence[BitSet], Sequence[int]]) – the scopes to evaluate, given either as a sequence of variable ids or
BitSetscopesforward (bool) – whether to enable partial evaluation for the forward pass
backward (bool) – whether to enable partial evaluation for the backward pass
overwrite (bool) – whether to overwrite an already-enabled partial-evaluation configuration
- TensorCircuit.disable_partial_evaluation(forward: bool = True, backward: bool = True)
Disable partial evaluation (see
enable_partial_evaluation()), so that subsequent passes again evaluate the entire circuit.