wip: CategoricalMassEstimator can be build and MixedDensityEstimator too. log_prob has shape issues tho

Author: jnsbck

sbi/made_mnle.ipynb

@@ -29,15 +29,16 @@ def __init__(
         use_batch_norm=False,
         epsilon=1e-2,
         custom_initialization=True,
+        #TODO: embedding_net: Optional[nn.Module] = None,
     ):
 
         if use_residual_blocks and random_mask:
             raise ValueError("Residual blocks can't be used with random masks.")
 
         self.num_variables = len(categories)
-        self.max_categories = max(categories)
+        self.num_categories = int(max(categories))
         self.categories = categories
-        self.mask = torch.zeros(self.num_variables, self.max_categories)
+        self.mask = torch.zeros(self.num_variables, self.num_categories)
         for i, c in enumerate(categories):
             self.mask[i, :c] = 1
 
@@ -46,7 +47,7 @@ def __init__(
             hidden_features,
             context_features=context_features,
             num_blocks=num_blocks,
-            output_multiplier=self.max_categories,
+            output_multiplier=self.num_categories,
             use_residual_blocks=use_residual_blocks,
             random_mask=random_mask,
             activation=activation,
@@ -68,10 +69,10 @@ def compute_probs(self, outputs):
         ps = ps / ps.sum(dim=-1, keepdim=True)
         return ps
 
-    # outputs (batch_size, num_variables, max_categories)
+    # outputs (batch_size, num_variables, num_categories)
     def log_prob(self, inputs, context=None):
         outputs = self.forward(inputs, context=context)
-        outputs = outputs.reshape(*inputs.shape, self.max_categories)
+        outputs = outputs.reshape(*inputs.shape, self.num_categories)
         ps = self.compute_probs(outputs)
 
         # categorical log prob
@@ -91,7 +92,7 @@ def sample(self, num_samples, context=None):
 
             for variable in range(self.num_variables):
                 outputs = self.forward(samples, context)
-                outputs = outputs.reshape(*samples.shape, self.max_categories)
+                outputs = outputs.reshape(*samples.shape, self.num_categories)
                 ps = self.compute_probs(outputs)
                 samples[:, variable] = Categorical(probs=ps[:,variable]).sample()
 

sbi/neural_nets/estimators/categorical_net.py

@@ -3,9 +3,9 @@
 
 from typing import Optional
 
-from torch import Tensor, nn, unique
+from torch import Tensor, nn, unique, tensor
 
-from sbi.neural_nets.estimators import CategoricalMassEstimator, CategoricalNet
+from sbi.neural_nets.estimators import CategoricalMassEstimator, CategoricalNet, CategoricalMADE
 from sbi.utils.nn_utils import get_numel
 from sbi.utils.sbiutils import (
     standardizing_net,
@@ -61,3 +61,53 @@ def build_categoricalmassestimator(
     return CategoricalMassEstimator(
         categorical_net, input_shape=batch_x[0].shape, condition_shape=batch_y[0].shape
     )
+
+
+def build_autoregressive_categoricalmassestimator(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    z_score_x: Optional[str] = "none",
+    z_score_y: Optional[str] = "independent",
+    num_hidden: int = 20,
+    num_layers: int = 2,
+    embedding_net: nn.Module = nn.Identity(),
+):
+    """Returns a density estimator for a categorical random variable.
+
+    Args:
+        batch_x: A batch of input data.
+        batch_y: A batch of condition data.
+        z_score_x: Whether to z-score the input data.
+        z_score_y: Whether to z-score the condition data.
+        num_hidden: Number of hidden units per layer.
+        num_layers: Number of hidden layers.
+        embedding_net: Embedding net for y.
+    """
+
+    if z_score_x != "none":
+        raise ValueError("Categorical input should not be z-scored.")
+
+    check_data_device(batch_x, batch_y)
+
+    z_score_y_bool, structured_y = z_score_parser(z_score_y)
+    y_numel = get_numel(batch_y, embedding_net=embedding_net)
+
+    if z_score_y_bool:
+        embedding_net = nn.Sequential(
+            standardizing_net(batch_y, structured_y), embedding_net
+        )
+
+
+    categories = tensor([unique(variable).numel() for variable in batch_x.T])
+
+    categorical_net = CategoricalMADE(
+        categories=categories,
+        context_features=y_numel,
+        hidden_features=num_hidden,
+        num_blocks=num_layers,
+        #TODO: embedding_net=embedding_net,
+    )
+
+    return CategoricalMassEstimator(
+        categorical_net, input_shape=batch_x[0].shape, condition_shape=batch_y[0].shape
+    )