added indexed categorical layer evaluation

cirkit/backend/torch/layers/input.py

@@ -296,6 +296,9 @@
                 raise ValueError(f"The number of folds and shape of 'probs' must match the layer's")
         self.probs = probs
         self.logits = logits
+        self.idx_mode = (
+            len(torch.unique(self.scope_idx)) > 4096 or self.num_categories > 256
+        )
 
     def _valid_parameter_shape(self, p: TorchParameter) -> bool:
         if p.num_folds != self.num_folds:
@@ -323,10 +326,30 @@
     def log_unnormalized_likelihood(self, x: Tensor) -> Tensor:
         if x.is_floating_point():
             x = x.long()  # The input to Categorical should be discrete
-        x = F.one_hot(x, self.num_categories)  # (F, C, B, 1, num_categories)
-        x = x.squeeze(dim=3)  # (F, C, B, num_categories)
         logits = torch.log(self.probs()) if self.logits is None else self.logits()
-        x = torch.einsum("fcbi,fkci->fbk", x.to(logits.dtype), logits)
+        if self.idx_mode:
+            if self.num_channels == 1:
+                x = (
+                    logits[:, :, 0, :]
+                    .transpose(1, 2)[range(self.num_folds), x[:, 0, :, 0].t()]
+                    .transpose(0, 1)
+                )
+            else:
+                x = x[..., 0].permute(2, 0, 1)
+                x = (
+                    logits[
+                        torch.arange(self.num_folds).unsqueeze(1),
+                        :,
+                        torch.arange(self.num_channels).unsqueeze(0),
+                        x,
+                    ]
+                    .sum(2)
+                    .transpose(0, 1)
+                )
+        else:
+            x = F.one_hot(x, self.num_categories)  # (F, C, B, 1, num_categories)
+            x = x.squeeze(dim=3)  # (F, C, B, num_categories)
+            x = torch.einsum("fcbi,fkci->fbk", x.to(logits.dtype), logits)
         return x
 
     def log_partition_function(self) -> Tensor: