Repository: april-tools/cirkit

Issue:
Sampling does not work in main branch?
Sampling does not seem to work with latest changes in the main branch. Gives the error [1].

```
samples = sampling_query(num_samples=sample_count)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "cirkit/cirkit/backend/torch/queries.py", line 234, in __call__
    samples = self._circuit.evaluate(
              ^^^^^^^^^^^^^^^^^^^^^^^
  File "cirkit/cirkit/backend/torch/graph/modules.py", line 295, in evaluate
    for module, inputs in self._address_book.lookup(module_outputs, in_graph=x):
  File "/cirkit/cirkit/backend/torch/circuits.py", line 58, in lookup
    x = in_graph[..., layer.scope_idx].permute(2, 1, 0, 3)
        ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^
TypeError: 'NoneType' object is not subscriptable
```

EDIT: Not sure if this is only with the latest changes, or if it's an already known issue, as I updated the code to the HEAD of the main after some time.

Hints from issue comments:
@loreloc: Hi Isuru, thanks a lot for reporting this.
I must have broken something in some of the last PRs.
I will fix it and add a unit test for the sampling.@isurulucky: Thanks a lot @loreloc !@loreloc: I should have fixed this in #332 . Let me know if you have other issues.@isurulucky: > I should have fixed this in #332 . Let me know if you have other issues.

I can confirm that this works. Thanks again @loreloc for the quick fix. 
A related question: are there plans to support conditional sampling, sampling with a marginalized circuit etc., in near future?@loreloc: Yes, there are plans to support those. Perhaps they will be available in 0.3.0.@isurulucky: > Yes, there are plans to support those. Perhaps they will be available in 0.3.0.

Great! thank you. 

Solution patch:
diff --git a/cirkit/backend/torch/circuits.py b/cirkit/backend/torch/circuits.py
index f1a283f..4fa91cb 100644
--- a/cirkit/backend/torch/circuits.py
+++ b/cirkit/backend/torch/circuits.py
@@ -48,15 +48,19 @@ class LayerAddressBook(AddressBook):
             # That is, we are gathering the inputs of input layers
             layer = entry.module
             assert isinstance(layer, TorchInputLayer)
-            if not layer.num_variables:
-                # Pass the wanted batch dimension to constant layers
-                yield layer, (1 if in_graph is None else in_graph.shape[0],)
-            else:
+            if layer.num_variables:
+                if in_graph is None:
+                    yield layer, ()
+                    continue
                 # in_graph: An input batch (assignments to variables) of shape (B, C, D)
                 # scope_idx: The scope of the layers in each fold, a tensor of shape (F, D'), D' < D
                 # x: (B, C, D) -> (B, C, F, D') -> (F, C, B, D')
                 x = in_graph[..., layer.scope_idx].permute(2, 1, 0, 3)
                 yield layer, (x,)
+                continue
+
+            # Pass the wanted batch dimension to constant layers
+            yield layer, (1 if in_graph is None else in_graph.shape[0],)
 
     @classmethod
     def from_index_info(


Test patch:
diff --git a/tests/backend/torch/test_compile_circuit.py b/tests/backend/torch/test_compile_circuit.py
index abe46f4..0f6065e 100644
--- a/tests/backend/torch/test_compile_circuit.py
+++ b/tests/backend/torch/test_compile_circuit.py
@@ -73,7 +73,7 @@ def check_continuous_ground_truth(
     assert isclose(int_tc(), semiring.map_from(torch.tensor(gt_partition_func), SumProductSemiring))
     df = lambda y, x: torch.exp(tc(torch.Tensor([[[x, y]]]))).squeeze()
     int_a, int_b = -np.inf, np.inf
-    ig, err = integrate.dblquad(df, int_a, int_b, int_a, int_b)
+    ig, err = integrate.dblquad(df, int_a, int_b, int_a, int_b, epsabs=1e-5, epsrel=1e-5)
     assert isclose(ig, gt_partition_func)
 
 
@@ -119,7 +119,6 @@ def test_compile_monotonic_structured_categorical_pc(fold: bool, optimize: bool,
     check_discrete_ground_truth(tc, int_tc, compiler.semiring, gt_outputs["evi"], gt_partition_func)
 
 
-@pytest.mark.slow
 def test_compile_monotonic_structured_gaussian_pc():
     compiler = TorchCompiler(fold=True, optimize=True, semiring="lse-sum")
     sc, gt_outputs, gt_partition_func = build_monotonic_bivariate_gaussian_hadamard_dense_pc(
diff --git a/tests/backend/torch/test_compile_circuit_operators.py b/tests/backend/torch/test_compile_circuit_operators.py
index 928528e..d56a0ce 100644
--- a/tests/backend/torch/test_compile_circuit_operators.py
+++ b/tests/backend/torch/test_compile_circuit_operators.py
@@ -105,7 +105,6 @@ def test_compile_product_integrate_pc_categorical(
     assert allclose(compiler.semiring.prod(each_tc_scores, dim=0), scores)
 
 
-@pytest.mark.slow
 def test_compile_product_integrate_pc_gaussian():
     compiler = TorchCompiler(semiring="lse-sum", fold=True, optimize=True)
     scs, tcs = [], []
@@ -113,7 +112,7 @@ def test_compile_product_integrate_pc_gaussian():
     num_products = 3
     for i in range(num_products):
         sci = build_bivariate_monotonic_structured_cpt_pc(
-            num_units=2 + i, input_layer="gaussian", normalized=False
+            num_units=1 + i, input_layer="gaussian", normalized=False
         )
         tci = compiler.compile(sci)
         scs.append(sci)
@@ -138,8 +137,8 @@ def test_compile_product_integrate_pc_gaussian():
     z = z.squeeze()
     df = lambda y, x: torch.exp(tc(torch.Tensor([[[x, y]]]))).squeeze()
     int_a, int_b = -np.inf, np.inf
-    ig, err = integrate.dblquad(df, int_a, int_b, int_a, int_b)
-    assert np.isclose(ig, torch.exp(z).item(), atol=1e-15)
+    ig, err = integrate.dblquad(df, int_a, int_b, int_a, int_b, epsabs=1e-5, epsrel=1e-5)
+    assert isclose(ig, torch.exp(z).item())
 
 
 @pytest.mark.parametrize(
diff --git a/tests/backend/torch/test_compile_marginalization.py b/tests/backend/torch/test_compile_marginalization.py
index 96838ba..c4eb43a 100644
--- a/tests/backend/torch/test_compile_marginalization.py
+++ b/tests/backend/torch/test_compile_marginalization.py
@@ -63,7 +63,6 @@ def test_marginalize_monotonic_pc_categorical(semiring: str, fold: bool, optimiz
         ), f"Input: {x}"
 
 
-@pytest.mark.slow
 def test_marginalize_monotonic_pc_gaussian():
     compiler = TorchCompiler(fold=True, optimize=True, semiring="lse-sum")
     sc, gt_outputs, gt_partition_func = build_monotonic_bivariate_gaussian_hadamard_dense_pc(
diff --git a/tests/backend/torch/test_queries/test_sampling.py b/tests/backend/torch/test_queries/test_sampling.py
index e69de29..a2a907a 100644
--- a/tests/backend/torch/test_queries/test_sampling.py
+++ b/tests/backend/torch/test_queries/test_sampling.py
@@ -0,0 +1,50 @@
+import itertools
+
+import pytest
+import torch
+
+from cirkit.backend.torch.circuits import TorchCircuit
+from cirkit.backend.torch.compiler import TorchCompiler
+from cirkit.backend.torch.queries import SamplingQuery
+from tests.floats import allclose
+from tests.symbolic.test_utils import build_multivariate_monotonic_structured_cpt_pc
+
+
+@pytest.mark.parametrize(
+    "fold,optimize",
+    itertools.product([False, True], [False, True]),
+)
+def test_quary_unconditional_sampling(fold: bool, optimize: bool):
+    compiler = TorchCompiler(semiring="lse-sum", fold=fold, optimize=optimize)
+    sc = build_multivariate_monotonic_structured_cpt_pc(
+        num_units=2, input_layer="bernoulli", parameterize=True, normalized=True
+    )
+    tc: TorchCircuit = compiler.compile(sc)
+
+    # Compute the probabilities
+    worlds = torch.tensor(list(itertools.product([0, 1], repeat=tc.num_variables))).unsqueeze(
+        dim=-2
+    )
+    assert worlds.shape == (2**tc.num_variables, 1, tc.num_variables)
+    tc_outputs = tc(worlds)
+    assert tc_outputs.shape == (2**tc.num_variables, 1, 1)
+    assert torch.all(torch.isfinite(tc_outputs))
+    probs = torch.exp(tc_outputs)
+    probs = probs.squeeze(dim=2).squeeze(dim=1)
+
+    # Sample data points unconditionally
+    num_samples = 1_000_000
+    query = SamplingQuery(tc)
+    # samples: (num_samples, C, D)
+    samples, _ = query(num_samples=num_samples)
+    assert samples.shape == (num_samples, 1, tc.num_variables)
+    samples = samples.squeeze(dim=1)
+
+    # Map samples to indices of the probabilities computed above
+    samples_idx = samples * torch.tensor(list(reversed([2**i for i in range(tc.num_variables)])))
+    samples_idx = torch.sum(samples_idx, dim=-1)
+
+    # Compute ratios and compare with the probabilities
+    _, counts = torch.unique(samples_idx, return_counts=True)
+    ratios = counts / num_samples
+    assert allclose(ratios, probs, atol=1e-3)