Remove lazy build from graph model

Author: pietrobarbiero

examples/utilization/1_pgm/0_concept_bottleneck_model.py

@@ -32,7 +32,7 @@ def main():
     # ParametricCPD setup
     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))
+    y_predictor = ParametricCPD("xor", parametrization=LinearCC(in_features_endogenous=2, out_features=2))
 
     # ProbabilisticModel Initialization
     concept_model = ProbabilisticModel(variables=[input_var, *concepts, tasks], parametric_cpds=[backbone, *c_encoder, y_predictor])

examples/utilization/2_model/0_concept_bottleneck_model.py

@@ -38,8 +38,8 @@ def main():
     concept_model = BipartiteModel(task_names,
                                    latent_dims,
                                    annotations,
-                                   LazyConstructor(LinearZC),
-                                   LazyConstructor(LinearCC))
+                                   LinearZC(10, 1),
+                                   LinearCC(2, 2))
 
     # Inference Initialization
     inference_engine = DeterministicInference(concept_model.probabilistic_model)

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

@@ -2,6 +2,7 @@
 
 import pandas as pd
 import torch
+from torch.nn import Module
 
 from .....annotations import Annotations
 from .concept_graph import ConceptGraph
@@ -35,7 +36,7 @@ class BipartiteModel(GraphModel):
     Example:
         >>> import torch
         >>> from torch_concepts import Annotations, AxisAnnotation
-        >>> from torch_concepts.nn import BipartiteModel, LazyConstructor
+        >>> from torch_concepts.nn import BipartiteModel, LazyConstructor, LinearCC
         >>> from torch.distributions import Bernoulli
         >>>
         >>> # Define concepts and tasks
@@ -78,11 +79,11 @@ def __init__(
             task_names: Union[List[str], str],
             input_size: int,
             annotations: Annotations,
-            encoder: LazyConstructor,
-            predictor: LazyConstructor,
+            encoder: Union[LazyConstructor, Module],
+            predictor: Union[LazyConstructor, Module],
             use_source_exogenous: bool = None,
-            source_exogenous: Optional[LazyConstructor] = None,
-            internal_exogenous: Optional[LazyConstructor] = None,
+            source_exogenous: Optional[Union[LazyConstructor, Module]] = None,
+            internal_exogenous: Optional[Union[LazyConstructor, Module]] = None,
     ):
         task_names = ensure_list(task_names)
         # get label names

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

@@ -1,5 +1,5 @@
-from typing import List, Tuple, Optional
-from torch.nn import Identity
+from typing import List, Tuple, Optional, Union
+from torch.nn import Identity, Module
 
 from .....annotations import Annotations
 from ..models.variable import Variable, InputVariable, ExogenousVariable, EndogenousVariable
@@ -50,7 +50,7 @@ class GraphModel(BaseConstructor):
         >>> import torch
         >>> import pandas as pd
         >>> from torch_concepts import Annotations, AxisAnnotation, ConceptGraph
-        >>> from torch_concepts.nn import GraphModel, LazyConstructor
+        >>> from torch_concepts.nn import GraphModel, LazyConstructor, LinearCC
         >>> from torch.distributions import Bernoulli
         >>>
         >>> # Define concepts and their structure
@@ -111,11 +111,11 @@ def __init__(self,
                  model_graph: ConceptGraph,
                  input_size: int,
                  annotations: Annotations,
-                 encoder: LazyConstructor,
-                 predictor: LazyConstructor,
+                 encoder: Union[LazyConstructor, Module],
+                 predictor: Union[LazyConstructor, Module],
                  use_source_exogenous: bool = None,
-                 source_exogenous: Optional[LazyConstructor] = None,
-                 internal_exogenous: Optional[LazyConstructor] = None
+                 source_exogenous: Optional[Union[LazyConstructor, Module]] = None,
+                 internal_exogenous: Optional[Union[LazyConstructor, Module]] = None
     ):
         super(GraphModel, self).__init__(
             input_size=input_size,
@@ -281,17 +281,7 @@ def _init_predictors(self,
                                      parents=endogenous_parents_names+exog_vars_names,
                                     distribution=self.annotations[1].metadata[c_name]['distribution'],
                                     size=self.annotations[1].cardinalities[self.annotations[1].get_index(c_name)])
-
-            # TODO: we currently assume predictors can use exogenous vars if any, but not latent
-            lazy_constructor = layer.build(
-                in_features_endogenous=in_features_endogenous,
-                in_features_exogenous=in_features_exogenous,
-                in_features=None,
-                out_features=predictor_var.size,
-                cardinalities=[predictor_var.size]
-            )
-
-            predictor_cpd = ParametricCPD(c_name, parametrization=lazy_constructor)
+            predictor_cpd = ParametricCPD(c_name, parametrization=layer)
 
             predictor_vars.append(predictor_var)
             predictor_cpds.append(predictor_cpd)