added backward sampling

[gengala]

supports TorchSumLayer, TorchCPTLayer, TorchHadamardLayer, and TorchCategoricalLayer input layers, on folded circuits. Other layers raise NotImplementedError. Unfolded circuits aren't supported yet.

For testing, train a circuit on mnist and then do:

import torch
import matplotlib.pyplot as plt

from cirkit.backend.torch.queries import SamplingQuery

circuit = torch.load("circuit.pt", weights_only=False)

# %% Top-down (backward) sampling 

torch.manual_seed(131)
sq_backward = SamplingQuery(circuit, backward=True)
td_samples, _ = sq_backward(num_samples=16)

images = td_samples.cpu().reshape(16, 28, 28).numpy()
fig, axes = plt.subplots(4, 4, figsize=(8, 8))
for i, ax in enumerate(axes.flat):
    ax.imshow(images[i], cmap="gray", vmin=0, vmax=255)
    ax.axis("off")
fig.suptitle("Top-down sampling (backward=True)")
plt.tight_layout()
plt.savefig("sampled_digits_integrated.png", dpi=100)
plt.show()

# %% Forward sampling 

sq_forward = SamplingQuery(circuit)

# NOTE: the forward sampler allocates (F, K, num_samples, D) tensors internally,
# which grow very large. On this circuit (F=784, K=64, D=784) even 16 samples
# can OOM. We sample 4 at a time and concatenate.
torch.manual_seed(42)
bu_batches = []
for _ in range(4):
    s, _ = sq_forward(num_samples=4)
    bu_batches.append(s.cpu())
bu_samples = torch.cat(bu_batches, dim=0)

images = bu_samples.reshape(16, 28, 28).numpy()
fig, axes = plt.subplots(4, 4, figsize=(8, 8))
for i, ax in enumerate(axes.flat):
    ax.imshow(images[i], cmap="gray", vmin=0, vmax=255)
    ax.axis("off")
fig.suptitle("Forward (bottom-up) sampling (backward=False)")
plt.tight_layout()
plt.savefig("sampled_digits_forward_integrated.png", dpi=100)
plt.show()

[codecov]

Codecov Report

❌ Patch coverage is 11.11111% with 120 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.44%. Comparing base (eb99f6f) to head (df52867).
⚠️ Report is 24 commits behind head on main.

Files with missing lines	Patch %	Lines
cirkit/backend/torch/queries.py	11.11%	119 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #403      +/-   ##
==========================================
- Coverage   68.71%   67.44%   -1.28%     
==========================================
  Files          58       58              
  Lines        6007     6140     +133     
  Branches      773      804      +31     
==========================================
+ Hits         4128     4141      +13     
- Misses       1590     1709     +119     
- Partials      289      290       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[chbricout]

Nice PR and a very needed speedup! My only comment is that I think the sampling logic should be implemented at the layer level like it's the case with forward sampling, otherwise it gets harder to create new layers because you need to know every query that needs new custom logic.

[arranger1044]

super cool, thanks @gengala!
I agree with @chbricout that we would need this to adopt the cirkit philosophy, could you please take care of it?
Also having benchmarks wrt past forward implementation would be nice.

[adrianjav]

I had a quick look and I agree with @chbricout and @arranger1044. On top of that, it worries me a bit the integration with the conditional-sampling branch. Now that the PR is monolithic it should not be a problem, but it could potentially be if we integrate forward and backward sampling in the layers.

Also having benchmarks wrt past forward implementation would be nice.

I have a local branch where I added pytest-benchmark and pytest-memray to use the unit tests directly as benchmarking. It should easy to merge, but it'd be great to have some agreed guidelines on benchmarking.

PS: I tested the PR wrt. the sampling tests from the conditional sampling branch, and the code passed all of those that implements (categorical leaves and folded circuits). They are non-exhaustive though.

[chbricout]

@gengala any news on this PR?