[hxu129]

Small irregular circuits can pay for CPU degree sorting, coalescing and Numba metadata compilation during juice.compile. This adds a bounded histogram path for dp_with_coalesce, then runs the existing weighted partition recurrence with NumPy integer tables.

The path accepts one-dimensional CPU integer tensors with positive degrees at most 4096, at most 64 occupied bins, 2–16 requested partitions, and max_degree * num_nodes < 2**31 - 1. It preserves the coalescing endpoints/counts, first-minimum ties, stopping threshold, return dtype, and existing fixed-column backtrace behavior. Two-bin histograms retain the existing weighted DP; inputs outside the bounded domain retain the existing dispatch. Circuit floating-point arithmetic is unchanged.

The focused tests include an independent scalar recurrence, comparisons with the unchanged compiled DP, dtype/noncontiguous-input and fallback cases, and an ordinary two-variable categorical circuit with irregular product-parent degrees. The CUDA cases compare compiled index tables and parameters, forward masses, linear/logspace leaf flows, and sample frequencies.

On an RTX 4090 with Python 3.10.14, PyTorch 2.4.0+cu121, Triton 3.0.0, NumPy 1.26.4 and Numba 0.66.0, the standalone benchmark compared unmodified upstream 3b44e329 with this patch. Each row uses five fresh processes per version on the same GPU, with OMP_NUM_THREADS=1. Times below are medians in seconds.

GPU compilation option	juice.compile, before → after	First request after runtime initialization, before → after	Including runtime initialization, before → after
Off	1.912 → 1.178	6.276 → 5.561	7.940 → 7.234
On	1.912 → 1.180	6.237 → 5.587	7.880 → 7.230

The first request includes PyJuice import, node construction, compilation, device transfer, and one native forward/backward query with CPU outputs for all nine observations. All 20 runs passed independent FP64 literal-mixture checks and produced identical recorded outputs across versions. PyTorch peak allocated CUDA tensor memory was 23,040 bytes in every run. These measurements describe this small fixture and software stack; disk caches were not cleared, and the timings are observations rather than confidence intervals. Other dependency versions and the full test suite have not been validated here.

To run the focused tests:

PYTHONPATH=src python -m pytest -q tests/layer/compact_histogram_partition_test.py

Result: 32 passed in 12.70 s: 30 CPU partition cases and two actual CUDA mixture cases, covering linear/logspace flows. Each CUDA case also checks the sample law with 8,192 draws. Against unmodified upstream, the same test file run with -k test_random_histograms_match_independent_recurrence_without_metadata_jit -x stops at the first selected case: the old path invokes _coalesce_fn and raises the test's injected AssertionError("coalesce JIT") (1 failed, 17 deselected, 3.81 s). This verifies that the regression check detects the metadata JIT being avoided.

To reproduce the old dispatch failure in a clean checkout of this PR, keep the new tests and restore only the partition implementation:

git show 3b44e329:src/pyjuice/layer/backend/node_partition.py > src/pyjuice/layer/backend/node_partition.py
PYTHONPATH=src python -m pytest -q tests/layer/compact_histogram_partition_test.py -k test_random_histograms_match_independent_recurrence_without_metadata_jit -x
git restore -- src/pyjuice/layer/backend/node_partition.py

To reproduce the timing boundary from this PR checkout, use its benchmark script for both source versions:

git worktree add --detach ../pyjuice-before 3b44e329b34e66e4f8ca7bb49b722b162a9607aa
for i in 1 2 3 4 5; do
  OMP_NUM_THREADS=1 CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 \
    PYTHONPATH=../pyjuice-before/src python examples/benchmark_native_compile.py \
    --no-gpu-compilation --output before-$i.json
  OMP_NUM_THREADS=1 CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 \
    PYTHONPATH=src python examples/benchmark_native_compile.py \
    --no-gpu-compilation --output after-$i.json
done

Repeat with --gpu-compilation and fresh output names. The JSON records individual stages, numerical checks, memory, versions and source fingerprints. No downloaded model or dataset is needed.

AI assistance was used to prepare the implementation, tests and this description. Validation and its limits are stated above.

[liuanji]

Closing, with thanks: the cost you found is real, and it is now addressed another way. These functions had no Numba cache, so every new process recompiled them. 41eb5f6 on main adds @njit(cache = True) to them, as elsewhere in the repo, with no second implementation of the DP.

juice.compile here, median of 4 fresh processes, warm cache:

Circuit	Before	This PR	Cached
This PR's example	1.109 s	0.881 s	0.909 s
HCLT-32	1.424 s	1.371 s	1.252 s

Your note on the fixed-column backtrace was useful as well: we will look at it when we revisit compilation.