Fix examples in docstrings under nn

Author: pietrobarbiero

torch_concepts/distributions/delta.py

@@ -33,6 +33,8 @@ class Delta(Distribution):
         mean: Returns the deterministic value.
 
     Examples:
+        >>> import torch
+        >>> from torch_concepts.distributions import Delta
         >>> dist = Delta(torch.tensor([1.0, 2.0, 3.0]))
         >>> sample = dist.sample()
         >>> print(sample)  # tensor([1., 2., 3.])

torch_concepts/nn/modules/high/base/learner.py

@@ -257,16 +257,6 @@ def configure_optimizers(self):
             Union[Optimizer, dict, None]: Returns optimizer directly, or dict with 
                 'optimizer' and optionally 'lr_scheduler' and 'monitor' keys,
                 or None if no optimizer is configured.
-                
-        Example:
-            >>> # With scheduler monitoring validation loss
-            >>> model = ConceptBottleneckModel(
-            ...     ...,
-            ...     optim_class=torch.optim.Adam,
-            ...     optim_kwargs={'lr': 0.001},
-            ...     scheduler_class=torch.optim.lr_scheduler.ReduceLROnPlateau,
-            ...     scheduler_kwargs={'mode': 'min', 'patience': 5, 'monitor': 'val_loss'}
-            ... )
         """
         # No optimizer configured
         if self.optim_class is None:

torch_concepts/nn/modules/low/dense_layers.py

@@ -45,6 +45,7 @@ def get_layer_activation(activation):
         ValueError: If activation name is not recognized.
         
     Example:
+        >>> from torch_concepts.nn.modules.low.dense_layers import get_layer_activation
         >>> act_class = get_layer_activation('relu')
         >>> activation = act_class()  # ReLU()
         >>> act_class = get_layer_activation(None)

torch_concepts/nn/modules/low/lazy.py

@@ -27,7 +27,7 @@ def _filter_kwargs_for_ctor(cls, **kwargs):
 
     Example:
         >>> import torch.nn as nn
-        >>> from torch_concepts.nn.modules.propagator import _filter_kwargs_for_ctor
+        >>> from torch_concepts.nn.modules.low.lazy import _filter_kwargs_for_ctor
         >>>
         >>> # Filter kwargs for Linear layer
         >>> kwargs = {'in_features': 10, 'out_features': 5, 'unknown_param': 42}
@@ -69,7 +69,7 @@ def instantiate_adaptive(module_cls, *args, drop_none=True, **kwargs):
 
     Example:
         >>> import torch.nn as nn
-        >>> from torch_concepts.nn.modules.propagator import instantiate_adaptive
+        >>> from torch_concepts.nn.modules.low.lazy import instantiate_adaptive
         >>>
         >>> # Instantiate a Linear layer with extra kwargs
         >>> kwargs = {'in_features': 10, 'out_features': 5, 'extra': None}
@@ -106,13 +106,13 @@ class LazyConstructor(torch.nn.Module):
         >>> from torch_concepts.nn import LinearCC
         >>>
         >>> # Create a propagator for a predictor
-        >>> lazy_constructorLazyConstructor(
+        >>> lazy_constructor = LazyConstructor(
         ...     LinearCC,
         ...     activation=torch.sigmoid
         ... )
         >>>
         >>> # Build the module when dimensions are known
-        >>> module = propagator.build(
+        >>> module = lazy_constructor.build(
         ...     out_features=3,
         ...     in_features_endogenous=5,
         ...     in_features=None,
@@ -121,7 +121,7 @@ class LazyConstructor(torch.nn.Module):
         >>>
         >>> # Use the module
         >>> x = torch.randn(2, 5)
-        >>> output = propagator(x)
+        >>> output = lazy_constructor(x)
         >>> print(output.shape)
         torch.Size([2, 3])
     """
@@ -227,12 +227,12 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
 
         Example:
             >>> import torch
-            >>> from torch_concepts.nn.modules.propagator import LazyConstructor
-            >>> from torch_concepts.nn.modules.predictors.linear import LinearCC
+            >>> from torch_concepts.nn import LazyConstructor
+            >>> from torch_concepts.nn import LinearCC
             >>>
             >>> # Create and build propagator
-            >>> lazy_constructorLazyConstructor(LinearCC)
-            >>> propagator.build(
+            >>> lazy_constructor = LazyConstructor(LinearCC)
+            >>> lazy_constructor.build(
             ...     out_features=3,
             ...     in_features_endogenous=5,
             ...     in_features=None,
@@ -241,7 +241,7 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
             >>>
             >>> # Forward pass
             >>> x = torch.randn(2, 5)
-            >>> output = propagator(x)
+            >>> output = lazy_constructor(x)
             >>> print(output.shape)
             torch.Size([2, 3])
         """

torch_concepts/nn/modules/metrics.py

@@ -36,18 +36,27 @@ class ConceptMetrics(nn.Module):
             specified concept names. Default: False.
             
     Example:
-        >>> from torch_concepts.nn.modules import ConceptMetrics, GroupConfig
+        >>> from torch_concepts.nn.modules.metrics import ConceptMetrics, GroupConfig
         >>> import torchmetrics
+        >>> import torch
+        >>> from torch_concepts import Annotations, AxisAnnotation
         >>> 
         >>> # Three ways to specify metrics:
+        >>> concept_annotations = Annotations({1: AxisAnnotation(
+        ...     labels=['concept1', 'concept2'],
+        ...     metadata={
+        ...         'concept1': {'type': 'discrete'},
+        ...         'concept2': {'type': 'discrete'}
+        ...     },
+        ... )})
         >>> metrics = ConceptMetrics(
         ...     annotations=concept_annotations,
         ...     fn_collection=GroupConfig(
         ...         binary={
         ...             # 1. Pre-instantiated
         ...             'accuracy': torchmetrics.classification.BinaryAccuracy(),
         ...             # 2. Class + user kwargs (average='macro')
-        ...             'f1': (torchmetrics.classification.BinaryF1Score, {'average': 'macro'})
+        ...             'f1': (torchmetrics.classification.BinaryF1Score, {'multidim_average': 'global'})
         ...         },
         ...         categorical={
         ...             # 3. Class only (num_classes will be added automatically)
@@ -57,7 +66,11 @@ class ConceptMetrics(nn.Module):
         ...     summary_metrics=True,
         ...     perconcept_metrics=['concept1', 'concept2']
         ... )
-        >>> 
+        >>>
+        >>> # Simulated predictions and targets
+        >>> predictions = torch.tensor([[0.8, 0.2], [0.4, 0.6]])
+        >>> targets = torch.tensor([[1, 0], [0, 1]])
+        >>>
         >>> # Update metrics during training
         >>> metrics.update(predictions, targets, split='train')
         >>> 

torch_concepts/nn/modules/mid/base/model.py

@@ -4,7 +4,10 @@
 This module provides the abstract base class for all concept-based models,
 defining the structure for models that use concept representations.
 """
+from typing import Union
+
 import torch
+from torch.nn import Module
 
 from .....annotations import Annotations
 from ...low.lazy import LazyConstructor
@@ -33,16 +36,26 @@ class BaseConstructor(torch.nn.Module):
     Example:
         >>> import torch
         >>> from torch_concepts import Annotations, AxisAnnotation
-        >>> from torch_concepts.nn import BaseModel, LazyConstructor
+        >>> from torch_concepts.nn import LazyConstructor
+        >>> from torch_concepts.nn.modules.mid.base.model import BaseConstructor
+        >>> from torch.distributions import RelaxedBernoulli
         >>>
         >>> # Create annotations for concepts
         >>> concept_labels = ('color', 'shape', 'size')
-        >>> annotations = Annotations({
-        ...     1: AxisAnnotation(labels=concept_labels)
-        ... })
+        >>> cardinalities = [1, 1, 1]
+        >>> metadata = {
+        ...     'color': {'distribution': RelaxedBernoulli},
+        ...     'shape': {'distribution': RelaxedBernoulli},
+        ...     'size': {'distribution': RelaxedBernoulli}
+        ... }
+        >>> annotations = Annotations({1: AxisAnnotation(
+        ...     labels=concept_labels,
+        ...     cardinalities=cardinalities,
+        ...     metadata=metadata
+        ... )})
         >>>
         >>> # Create a concrete model class
-        >>> class MyConceptModel(BaseModel):
+        >>> class MyConceptModel(BaseConstructor):
         ...     def __init__(self, input_size, annotations, encoder, predictor):
         ...         super().__init__(input_size, annotations, encoder, predictor)
         ...         # Build encoder and predictor
@@ -84,8 +97,8 @@ class BaseConstructor(torch.nn.Module):
     def __init__(self,
                  input_size: int,
                  annotations: Annotations,
-                 encoder: LazyConstructor,  # layer for root concepts
-                 predictor: LazyConstructor,
+                 encoder: Union[LazyConstructor, Module],  # layer for root concepts
+                 predictor: Union[LazyConstructor, Module],
                  *args,
                  **kwargs
     ):

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

@@ -158,6 +158,8 @@ def from_sparse(cls, edge_index: Tensor, edge_weight: Tensor, n_nodes: int, node
             ConceptGraph instance
             
         Example:
+            >>> import torch
+            >>> from torch_concepts import ConceptGraph
             >>> edge_index = torch.tensor([[0, 0, 1], [1, 2, 2]])
             >>> edge_weight = torch.tensor([1.0, 1.0, 1.0])
             >>> graph = ConceptGraph.from_sparse(edge_index, edge_weight, n_nodes=3)
@@ -310,11 +312,6 @@ def to_networkx(self, threshold: float = 0.0) -> nx.DiGraph:
 
         Returns:
             nx.DiGraph: NetworkX directed graph
-
-        Example:
-            >>> G = graph.to_networkx()
-            >>> list(G.nodes())
-            ['A', 'B', 'C']
         """
         # If threshold is 0.0 and we have a cache, return it
         if threshold == 0.0 and self._nx_graph_cache is not None:
@@ -357,11 +354,6 @@ def dense_to_sparse(self, threshold: float = 0.0) -> Tuple[Tensor, Tensor]:
         Returns:
             edge_index: Tensor of shape (2, num_edges) with source and target indices
             edge_weight: Tensor of shape (num_edges,) with edge weights
-
-        Example:
-            >>> edge_index, edge_weight = graph.dense_to_sparse()
-            >>> edge_index.shape
-            torch.Size([2, num_edges])
         """
         if threshold > 0.0:
             # Filter edges by threshold
@@ -500,6 +492,8 @@ def dense_to_sparse(
         edge_weight: Tensor of shape (num_edges,) with edge weights
 
     Example:
+        >>> import torch
+        >>> from torch_concepts.nn.modules.mid.constructors.concept_graph import dense_to_sparse
         >>> adj = torch.tensor([[0., 1., 0.],
         ...                     [0., 0., 1.],
         ...                     [0., 0., 0.]])
@@ -539,6 +533,8 @@ def to_networkx_graph(
         nx.DiGraph: NetworkX directed graph
 
     Example:
+        >>> import torch
+        >>> from torch_concepts.nn.modules.mid.constructors.concept_graph import to_networkx_graph
         >>> adj = torch.tensor([[0., 1., 1.],
         ...                     [0., 0., 1.],
         ...                     [0., 0., 0.]])
@@ -590,6 +586,8 @@ def get_root_nodes(
         List of root node names
 
     Example:
+        >>> import torch
+        >>> from torch_concepts.nn.modules.mid.constructors.concept_graph import get_root_nodes
         >>> adj = torch.tensor([[0., 1., 1.],
         ...                     [0., 0., 1.],
         ...                     [0., 0., 0.]])
@@ -621,6 +619,8 @@ def get_leaf_nodes(
         List of leaf node names
 
     Example:
+        >>> import torch
+        >>> from torch_concepts.nn.modules.mid.constructors.concept_graph import get_leaf_nodes
         >>> adj = torch.tensor([[0., 1., 1.],
         ...                     [0., 0., 1.],
         ...                     [0., 0., 0.]])
@@ -657,6 +657,8 @@ def topological_sort(
         nx.NetworkXError: If graph contains cycles
 
     Example:
+        >>> import torch
+        >>> from torch_concepts.nn.modules.mid.constructors.concept_graph import topological_sort
         >>> adj = torch.tensor([[0., 1., 1.],
         ...                     [0., 0., 1.],
         ...                     [0., 0., 0.]])
@@ -692,6 +694,8 @@ def get_predecessors(
         List of predecessor node names
 
     Example:
+        >>> import torch
+        >>> from torch_concepts.nn.modules.mid.constructors.concept_graph import get_predecessors
         >>> adj = torch.tensor([[0., 1., 1.],
         ...                     [0., 0., 1.],
         ...                     [0., 0., 0.]])
@@ -731,6 +735,8 @@ def get_successors(
         List of successor node names
 
     Example:
+        >>> import torch
+        >>> from torch_concepts.nn.modules.mid.constructors.concept_graph import get_successors
         >>> adj = torch.tensor([[0., 1., 1.],
         ...                     [0., 0., 1.],
         ...                     [0., 0., 0.]])

torch_concepts/nn/modules/mid/inference/forward.py

@@ -712,7 +712,7 @@ class DeterministicInference(ForwardInference):
         >>> from torch.distributions import Bernoulli
         >>> from torch_concepts import InputVariable, EndogenousVariable
         >>> from torch_concepts.distributions import Delta
-        >>> from torch_concepts.nn import DeterministicInference, ParametricCPD, ProbabilisticModel
+        >>> from torch_concepts.nn import DeterministicInference, ParametricCPD, ProbabilisticModel, LinearCC
         >>>
         >>> # Create a simple PGM: latent -> A -> B
         >>> input_var = InputVariable('input', parents=[], distribution=Delta, size=10)
@@ -723,7 +723,7 @@ class DeterministicInference(ForwardInference):
         >>> from torch.nn import Identity, Linear
         >>> cpd_emb = ParametricCPD('input', parametrization=Identity())
         >>> cpd_A = ParametricCPD('A', parametrization=Linear(10, 1))
-        >>> cpd_B = ParametricCPD('B', parametrization=Linear(1, 1))
+        >>> cpd_B = ParametricCPD('B', parametrization=LinearCC(1, 1))
         >>>
         >>> # Create probabilistic model
         >>> pgm = ProbabilisticModel(
@@ -743,7 +743,7 @@ class DeterministicInference(ForwardInference):
         >>> print(results['B'].shape)  # torch.Size([4, 1]) - endogenous, not {0,1}
         >>>
         >>> # Query specific concepts - returns concatenated endogenous
-        >>> output = inference.query(['B', 'A'], evidence={'embedding': x})
+        >>> output = inference.query(['B', 'A'], evidence={'input': x})
         >>> print(output.shape)  # torch.Size([4, 2])
         >>> # output contains [logit_B, logit_A] for each sample
         >>>
@@ -794,6 +794,8 @@ class AncestralSamplingInference(ForwardInference):
         >>> from torch_concepts import InputVariable
         >>> from torch_concepts.distributions import Delta
         >>> from torch_concepts.nn import AncestralSamplingInference, ParametricCPD, ProbabilisticModel
+        >>> from torch_concepts import EndogenousVariable
+        >>> from torch_concepts.nn import LinearCC
         >>>
         >>> # Create a simple PGM: embedding -> A -> B
         >>> embedding_var = InputVariable('embedding', parents=[], distribution=Delta, size=10)
@@ -804,7 +806,7 @@ class AncestralSamplingInference(ForwardInference):
         >>> from torch.nn import Identity, Linear
         >>> cpd_emb = ParametricCPD('embedding', parametrization=Identity())
         >>> cpd_A = ParametricCPD('A', parametrization=Linear(10, 1))
-        >>> cpd_B = ParametricCPD('B', parametrization=Linear(1, 1))
+        >>> cpd_B = ParametricCPD('B', parametrization=LinearCC(1, 1))
         >>>
         >>> # Create probabilistic model
         >>> pgm = ProbabilisticModel(
@@ -838,8 +840,8 @@ class AncestralSamplingInference(ForwardInference):
         >>>
         >>> # With relaxed distributions (requires temperature)
         >>> from torch.distributions import RelaxedBernoulli
-        >>> var_A_relaxed = Variable('A', parents=['embedding'],
-        ...                          distribution=RelaxedBernoulli, size=1)
+        >>> var_A_relaxed = InputVariable('A', parents=['embedding'],
+        ...                               distribution=RelaxedBernoulli, size=1)
         >>> pgm = ProbabilisticModel(
         ...     variables=[embedding_var, var_A_relaxed, var_B],
         ...     parametric_cpds=[cpd_emb, cpd_A, cpd_B]

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

@@ -60,9 +60,9 @@ class ParametricCPD(nn.Module):
     ...     parametrization=[module_a, module_b]
     ... )
     >>>
-    >>> print(cpd[0].module)
+    >>> print(cpd[0].parametrization)
     Linear(in_features=10, out_features=1, bias=True)
-    >>> print(cpd[1].module)
+    >>> print(cpd[1].parametrization)
     Sequential(...)
 
     Notes