Allow to instantiate parametricCPD with lazy constructors, otherwise assume the layer is already correctly and fully instantiated

Author: pietrobarbiero

examples/utilization/1_pgm/0_concept_bottleneck_model.py

@@ -5,7 +5,7 @@
 from torch_concepts import Annotations, AxisAnnotation, Variable, InputVariable, EndogenousVariable
 from torch_concepts.data.datasets import ToyDataset
 from torch_concepts.nn import LinearZC, LinearCC, ParametricCPD, ProbabilisticModel, \
-    RandomPolicy, DoIntervention, intervention, DeterministicInference
+    RandomPolicy, DoIntervention, intervention, DeterministicInference, LazyConstructor
 
 
 def main():
@@ -30,9 +30,9 @@ def main():
     tasks = EndogenousVariable("xor", parents=concept_names, distribution=RelaxedOneHotCategorical, size=2)
 
     # ParametricCPD setup
-    backbone = ParametricCPD("input", parametrization=torch.nn.Identity())
-    c_encoder = ParametricCPD(["c1", "c2"], parametrization=LinearZC(in_features=x_train.shape[1], out_features=concepts[0].size))
-    y_predictor = ParametricCPD("xor", parametrization=LinearCC(in_features_endogenous=sum(c.size for c in concepts), out_features=tasks.size))
+    backbone = ParametricCPD("input", parametrization=torch.nn.Sequential(torch.nn.Linear(x_train.shape[1], latent_dims), torch.nn.LeakyReLU()))
+    c_encoder = ParametricCPD(["c1", "c2"], parametrization=LazyConstructor(LinearZC))
+    y_predictor = ParametricCPD("xor", parametrization=LazyConstructor(LinearCC))
 
     # ProbabilisticModel Initialization
     concept_model = ProbabilisticModel(variables=[input_var, *concepts, tasks], parametric_cpds=[backbone, *c_encoder, y_predictor])

examples/utilization/1_pgm/1_concept_bottleneck_model_ancestral_sampling.py

@@ -5,7 +5,7 @@
 from torch_concepts import Annotations, AxisAnnotation, Variable, InputVariable, EndogenousVariable
 from torch_concepts.data.datasets import ToyDataset
 from torch_concepts.nn import LinearZC, LinearCC, ParametricCPD, ProbabilisticModel, \
-    RandomPolicy, DoIntervention, intervention, AncestralSamplingInference
+    RandomPolicy, DoIntervention, intervention, AncestralSamplingInference, LazyConstructor
 
 
 def main():
@@ -24,14 +24,14 @@ def main():
     y_train = torch.cat([y_train, 1-y_train], dim=1)
 
     # Variable setup
-    input_var = InputVariable("input", parents=[], size=latent_dims)
+    input_var = InputVariable("input", parents=[], size=x_train.shape[1])
     concepts = EndogenousVariable(concept_names, parents=["input"], distribution=RelaxedBernoulli)
     tasks = EndogenousVariable("xor", parents=concept_names, distribution=RelaxedOneHotCategorical, size=2)
 
     # ParametricCPD setup
     backbone = ParametricCPD("input", parametrization=torch.nn.Identity())
-    c_encoder = ParametricCPD(["c1", "c2"], parametrization=LinearZC(in_features=x_train.shape[1], out_features=concepts[0].size))
-    y_predictor = ParametricCPD("xor", parametrization=LinearCC(in_features_endogenous=sum(c.size for c in concepts), out_features=tasks.size))
+    c_encoder = ParametricCPD(["c1", "c2"], parametrization=LazyConstructor(LinearZC))
+    y_predictor = ParametricCPD("xor", parametrization=LazyConstructor(LinearCC))
 
     # ProbabilisticModel Initialization
     concept_model = ProbabilisticModel(variables=[input_var, *concepts, tasks], parametric_cpds=[backbone, *c_encoder, y_predictor])

examples/utilization/2_model/4_concept_graph_model_learned.py

@@ -57,7 +57,7 @@ def main():
                                    source_exogenous=LazyConstructor(LinearZU, exogenous_size=11),
                                    internal_exogenous=LazyConstructor(LinearZU, exogenous_size=7),
                                    encoder=LazyConstructor(LinearUC),
-                                   predictor=LazyConstructor(HyperLinearCUC, embedding_size=20),)
+                                   predictor=LazyConstructor(HyperLinearCUC, embedding_size=20))
 
     # graph learning init
     graph_learner = WANDAGraphLearner(concept_names, task_names)

torch_concepts/nn/modules/mid/constructors/graph.py

@@ -188,14 +188,7 @@ def _init_exog(self, layer: LazyConstructor, label_names, parent_var, cardinalit
                             distribution=Delta,
                             size=layer._module_kwargs['exogenous_size'])
 
-        lazy_constructor = layer.build(
-            in_features=parent_var.size,
-            in_features_endogenous=None,
-            in_features_exogenous=None,
-            out_features=1,
-        )
-
-        exog_cpds = ParametricCPD(exog_names, parametrization=lazy_constructor)
+        exog_cpds = ParametricCPD(exog_names, parametrization=layer)
         return exog_vars, exog_cpds
 
     def _init_encoder(self, layer: LazyConstructor, label_names, parent_vars, cardinalities=None) -> Tuple[Variable, ParametricCPD]:
@@ -220,13 +213,7 @@ def _init_encoder(self, layer: LazyConstructor, label_names, parent_vars, cardin
             if not isinstance(encoder_vars, list):
                 encoder_vars = [encoder_vars]
 
-            lazy_constructor = layer.build(
-                in_features=parent_vars[0].size,
-                in_features_endogenous=None,
-                in_features_exogenous=None,
-                out_features=encoder_vars[0].size,
-            )
-            encoder_cpds = ParametricCPD(label_names, parametrization=lazy_constructor)
+            encoder_cpds = ParametricCPD(label_names, parametrization=layer)
             # Ensure encoder_cpds is always a list
             if not isinstance(encoder_cpds, list):
                 encoder_cpds = [encoder_cpds]
@@ -241,13 +228,7 @@ def _init_encoder(self, layer: LazyConstructor, label_names, parent_vars, cardin
                                     parents=exog_vars_names,
                                     distribution=self.annotations[1].metadata[label_name]['distribution'],
                                     size=self.annotations[1].cardinalities[self.annotations[1].get_index(label_name)])
-                lazy_constructor = layer.build(
-                    in_features=None,
-                    in_features_endogenous=None,
-                    in_features_exogenous=exog_vars[0].size,
-                    out_features=encoder_var.size,
-                )
-                encoder_cpd = ParametricCPD(label_name, parametrization=lazy_constructor)
+                encoder_cpd = ParametricCPD(label_name, parametrization=layer)
                 encoder_vars.append(encoder_var)
                 encoder_cpds.append(encoder_cpd)
         return encoder_vars, encoder_cpds

torch_concepts/nn/modules/mid/models/probabilistic_model.py

@@ -9,7 +9,8 @@
 from torch.distributions import Distribution
 from typing import List, Dict, Optional, Type
 
-from .variable import Variable, ExogenousVariable
+from torch_concepts.nn import LazyConstructor
+from .variable import Variable, ExogenousVariable, EndogenousVariable, InputVariable
 from .cpd import ParametricCPD
 
 
@@ -159,14 +160,34 @@ def _initialize_model(self, input_parametric_cpds: List[ParametricCPD]):
                 if concept in self.concept_to_variable:
                     parametric_cpd.variable = self.concept_to_variable[concept]
                     parametric_cpd.parents = self.concept_to_variable[concept].parents
-                if not isinstance(parametric_cpd.variable, ExogenousVariable):
-                    new_parametrization = _reinitialize_with_new_param(parametric_cpd.parametrization,
-                                                                       'out_features',
-                                                                       self.concept_to_variable[concept].size)
-                    new_parametric_cpd = ParametricCPD(concepts=[concept], parametrization=new_parametrization)
-                    self.parametric_cpds[concept] = new_parametric_cpd
+
+                if isinstance(parametric_cpd.parametrization, LazyConstructor):
+                    parent_vars = [self.concept_to_variable[parent_ref] for parent_ref in parametric_cpd.variable.parents]
+                    in_features_endogenous = in_features_exogenous = in_features = 0
+                    for pv in parent_vars:
+                        if isinstance(pv, ExogenousVariable):
+                            in_features_exogenous = pv.size
+                        elif isinstance(pv, EndogenousVariable):
+                            in_features_endogenous += pv.size
+                        else:
+                            in_features += pv.size
+
+                    if isinstance(parametric_cpd.variable, ExogenousVariable):
+                        out_features = 1
+                    else:
+                        out_features = self.concept_to_variable[concept].size
+
+                    initialized_layer = parametric_cpd.parametrization.build(
+                        in_features=in_features,
+                        in_features_endogenous=in_features_endogenous,
+                        in_features_exogenous=in_features_exogenous,
+                        out_features=out_features,
+                    )
+                    new_parametrization = ParametricCPD(concepts=[concept], parametrization=initialized_layer)
                 else:
-                    self.parametric_cpds[concept] = parametric_cpd
+                    new_parametrization = parametric_cpd
+
+                self.parametric_cpds[concept] = new_parametrization
 
         # ---- Parent resolution (unchanged) ----
         for var in self.variables: