[liuanji]

Summary

This PR removes partial evaluation: the feature that restricted forward and backward passes to the nodes whose scope lies inside a given set of variables. It had not been used for a long time, and it had produced silent bugs:

• a CUDA graph replayed across a partial-evaluation toggle gave log-likelihoods off by up to 38 nats (Pc perf followups #58);
• input layers used block ids as node offsets (fix InputLayer partial_eval using block ids as per-node offsets #40);
• conditional(target_vars=...) was wrong for DiscreteLogistic leaves, as described below.

Its last internal user was queries.conditional's target_vars, which now builds its own node list. The removal makes PCs 0–4% faster in eager mode and changes no results.

Closes #21.

Bug fix: conditional(..., target_vars=...) with DiscreteLogistic leaves

conditional(pc, ..., target_vars = T) restricted its input-layer kernels to the target nodes through partial evaluation. The DiscreteLogistic kernel then read each target node's parameters and metadata through the node's position in the target list instead of its node id. So every node was combined with another node's mean, scale and value range. The error was up to 0.34 in probability, while every row still summed to 1. The Categorical path's version of this bug was #21, fixed by #40; the DiscreteLogistic path still had it.

The query now builds an explicit list of the nodes whose variable is a target. Both conditional kernels always index through it (for the full query it is simply all nodes), so the bug can't recur, and target_vars stays cheap for large vocabularies.

The new test tests/queries/conditional/target_vars_test.py checks target_vars against the full conditional:

• for Categorical and DiscreteLogistic leaves;
• on HMM, HCLT, PD and RAT-SPN;
• with sorted, unsorted and single targets, in both call orders.

The DiscreteLogistic cases fail on main and pass here.

Breaking changes

• TensorCircuit.enable_partial_evaluation / disable_partial_evaluation are removed, together with the same methods on LayerGroup, Layer, ProdLayer and InputLayer, and the scope-to-node-id caches that only fed them.
• ProdLayer.forward no longer takes _for_backward. It only selected the partial-evaluation nodes for the forward recompute inside the backward pass; that recompute is now the ordinary forward. A caller still passing it lands in **kwargs and gets the identical result.
• Custom distributions with post-processing kernels: the launch arguments the input layer hands them no longer include fw_local_ids_ptr / bk_local_ids_ptr or partial_eval, so such kernels must drop those parameters. Every built-in distribution is updated.

What changed internally

• 18 Triton kernels in layer/kernels/ (sum forward and backward, product) and the input-layer kernel templates lose their local-id pointer, their partial_eval compile-time constant, and the remap it guarded. The constant was always 0, so the compiled arithmetic is unchanged.
• The sum, product and input layers keep only their whole-layer code path. The CUDA fast-path gates lose their local_ids is None conditions, and the autotune cache keys lose an entry that was always zero.
• Distribution kernels updated: ExternProductCategorical (8), SoftEvidenceCategorical (6), LatentSoftEvidence, External, SoftEvidenceIndicator, plus an always-False partial flag in queries.conditional's soft-evidence forward kernels.
• Tests and docs: the partial-evaluation tests are removed. The target_vars case of input_layer_partial_eval_block_test.py is covered by the new test on the same kind of input layers (one block per input group). The API reference no longer lists the removed methods.

Verification

• Output equivalence: an eager differential harness compares 27 workloads (110 outputs) against main. These cover forward, backward, EM, MPE, tempering, sparse mode, HCLT, PD, RAT-SPN, soft evidence, BlockScale, LowRank, sampling, marginal / conditional, and every distribution touched here.

  • With autotuning off, every deterministic output is bit-identical.
  • The remaining differences are atomic-add noise within main's own run-to-run spread, or timing-based kernel choices that match main when rerun in isolation.

• Full test suite after each commit: 1252 passed, 26 skipped.

• Speed, eager forward+backward step, main → this PR, interleaved runs on the same GPU:

  Workload	Before	After
  HMM, 1024 latents, batch 64	3.78 ms	3.78 ms
  HCLT-256, batch 512	1.02 ms	0.98 ms (−3.6%)
  1-D PD, batch 256	2.72 ms	2.60 ms (−4.4%)
  Soft-evidence HMM, top-k, batch 16	3.71 ms	3.61 ms (−2.6%)

  Less per-launch host work; on the last three workloads every run of this PR was faster than every run of main.

• Training:

  • Standalone: HCLT-256 on ImageNet32 with mini-batch EM, at batch 512 and batch 64. Step times and training log-likelihoods are identical within run-to-run noise.
  • CoDD 2.0: one full epoch of LLaDA-8B LoRA joint training ran at 367 ms/step on both main and this PR, with validation log-likelihood −0.8074 vs −0.8077. Measured on its own, the PC's forward+backward takes 5.42 ms on main and 5.30 ms on this PR.

Commits

031aa59e fix(queries): conditional's target_vars no longer uses partial evaluation
328bf880 refactor!: remove the partial-evaluation API
36f55858 refactor!: remove partial evaluation from the sum and product layers and their kernels
2b5d651c refactor!: remove partial evaluation from the input layers and distributions