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 = 32, device: int | device | None = None, verbose: bool = True, cuda_graphs: bool = False, precision: str = 'auto')
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)
- note:
the default is high (32) because the memory/locality side of that trade dominates on current hardware – splitting tied groups across accumulators multiplies the flow buffer, and the extra traffic costs more than the atomic contention it avoids. Measured on a homogeneous HMM (seq 32, 1024 latents, 126464 cats, top-k 1024 soft evidence), going from 4 accumulator blocks to 1: param_flows 4.14 -> 1.04 GB, backward 16.8 -> 13.4 ms and init_param_flows 5.7 -> 1.4 ms (1.41x on the step), with parameters after an EM step agreeing to ~2e-5 relative – flow-accumulation reassociation only. For scale, on the same GPU 268M `atomic_add`s all targeting ONE address cost 3.7 ms, versus 44 ms for the same count scattered over a 4 GB table: contention is the cheap side of the trade. Lower it only for a model where atomic stalls are demonstrably the bottleneck.
verbose (bool) – Whether to display the progress of the compilation
cuda_graphs (bool) – OFF by default. When on, a forward or backward that runs a second time with the same buffers and options records its inner layers into a CUDA graph and later identical calls replay it, as if record_cudagraph = True were passed (an explicit record_cudagraph still wins). Worth it when the host bounds a step – small batches. Anything that cannot be captured runs eagerly.
Inference
- TensorCircuit.forward(inputs: Tensor, input_layer_fn: str | Callable | None = None, cache: dict | None = None, return_cache: bool = False, record_cudagraph: bool | None = None, apply_cudagraph: bool = True, 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
record_cudagraph (Optional[bool]) – True: capture the inner layers into a CUDA graph the first time this pass runs with these buffers and options (the call itself runs normally), so later identical calls replay it without the per-layer Python. Worth it when the host, not the GPU, bounds a step – small batches. Graphs are keyed on everything the inner layers are given (buffers by address, every option and keyword argument except the input layers’ per-call evidence, see Distribution.call_kwargs), at most max_cuda_graphs are kept, .to() drops them, and a pass that cannot be captured runs eagerly. None (the default): the circuit’s cuda_graphs setting, which is off unless it was compiled with it. False: do not record
apply_cudagraph (bool) – replay a graph recorded under exactly this call’s signature, if there is one
- 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 | None = None, apply_cudagraph: bool = True, allow_modify_flows: bool = False, propagation_alg: str | Sequence[str] = 'LL', logspace_flows: bool = True, negate_pflows: bool = False, _inner_layers_only: bool = False, _disable_buffer_init: bool = False, pflow_temperature: float = 1.0, temper_eflow: bool = False, compute_external_grads: bool = True, **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
record_cudagraph (Optional[bool]) – as in
forward(); a backward given layer callbacks always runs eagerly
- 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, _apply_denom_correction: bool = True)
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
_apply_denom_correction (bool) – internal. False runs the PLAIN M-step even on a circuit whose parameterization requests the denominator flow (
apply_z_correction). For callers that use this method as a RENORMALIZER rather than as an EM step – the gradient optimizers project back onto the normalized manifold withparam_flowsholding the partition flow and the denominator deliberately zeroed – the conditional update is not merely inappropriate, it divides by pseudocount * theta and returns inf atpseudocount = 0. Not for user code; an EM caller always wants the default.
- 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.