pyjuice.queries.sample

pyjuice.queries.sample(pc: TensorCircuit, num_samples: int | None = None, conditional: bool = False, use_cudagraph: bool = False, _use_scope_plan: bool = True, _sample_input_ns: bool = True, _do_calibration: bool = False, **kwargs)

Draw samples from a PC by performing a top-down ancestral sampling pass.

For unconditional sampling, num_samples must be specified. For conditional sampling (conditional = True), run a forward pass on the evidence first (e.g., via marginal()); the sampler then reuses the cached pc.node_mars and draws one sample per example in that batch, so num_samples is ignored.

Parameters:

• pc (TensorCircuit) – the input PC

• num_samples (Optional[int]) – number of samples to draw; required for unconditional sampling, ignored when conditional = True

• conditional (bool) – whether to sample conditioned on the evidence cached by a preceding forward pass

Per-sample external sum parameters are supplied through sum_external_params, in kwargs, exactly as they are to TensorCircuit.forward()

samples = juice.queries.sample(pc, num_samples = 1024, sum_external_params = {ns: phi})

with one gate per DRAWN SAMPLE, so their batch axis is num_samples. A sum layer that is given none is sampled from its shared parameters, which is what an ungated forward pass computes for it too. Under conditional = True they are instead taken from the forward pass that produced pc.node_mars – element_mars was built under those gates, so the draw has to use them – and passing them here only has to name the same nodes.

Parameters:

• use_cudagraph (bool) – capture the top-down pass into a CUDA graph and replay it. OPT-IN, because a graph owns a private memory pool and the frontier buffers it was captured with, held for the circuit’s lifetime – worth it for a sampling loop, not for a one-off draw. Needs pc.is_structured_decomposable (the index plan must repeat, or a replay is simply wrong) and is refused otherwise. One graph is kept per (num_samples, conditional), bounded with the plans.

• _sample_input_ns (bool) – whether to finish the draw by emitting a value from each selected input node. Default True, which is what makes the return a tensor of VALUES. Setting it False stops one step earlier and returns the FRONTIER instead – see the note below, which is the whole contract.

Returns:

with _sample_input_ns = True (the default), samples of size [num_samples, num_vars] in the input distributions’ own dtype.

With _sample_input_ns = False, the frontier: an int64 tensor of shape [rows, num_samples] holding the NODE ID each sample selected, padded with -1. rows is the frontier’s height, which is at least num_vars and usually more; each column is a dense prefix of live ids followed by nothing but -1, so >= 0 selects exactly the live entries, and every column of one draw holds exactly num_vars of them – one per variable, never two.

Return type:

torch.Tensor

Note:

the frontier’s ids are GLOBAL node ids, and a circuit may have more than one input layer – it has one per distribution TYPE, so a model mixing, say, Categorical and Bernoulli variables has two. Decoding therefore has to attribute each id to its owning layer BEFORE subtracting a start index:

for node_id in frontier[frontier[:, j] >= 0, j].tolist():
    for layer in pc.input_layer_group:
        start, end = layer._output_ind_range
        if start <= node_id < end:
            var = int(layer.vids[node_id - start, 0])
            params_start = int(layer.s_pids[node_id - start])
            break

The shortcut of subtracting pc.input_layer_group[0]._output_ind_range[0] and indexing that layer’s vids is correct only while the circuit has ONE input layer. It does not fail gracefully when it stops being true: ids from the second layer index past the end of the first layer’s vids, which is a device-side assert that poisons the CUDA context.