nous 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

nous/model.py

@@ -0,0 +1,226 @@
+from __future__ import annotations
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, List, Optional, Sequence, Tuple, Union
+
+from .facts import BetaFactLayer, PiecewiseLinearCalibrator
+from .prototypes import ScaledPrototypeLayer
+from .rules.blocks import SimpleNousBlock
+from .rules.softmax import SoftmaxRuleLayer
+from .rules.sparse import SparseRuleLayer
+
+class NousNet(nn.Module):
+    """
+    NousNet: rule-based neural network with honest interpretability for classification and regression.
+    """
+    def __init__(
+        self,
+        input_dim: int,
+        num_outputs: int,
+        task_type: str = "classification",     # "classification" | "regression"
+        feature_names: Optional[Sequence[str]] = None,
+        num_facts: int = 48,
+        rules_per_layer: Sequence[int] = (24, 12),
+        use_calibrators: bool = False,
+        rule_selection_method: str = "fixed",  # "fixed" | "softmax" | "sparse"
+        use_prototypes: bool = False,
+        l0_lambda: float = 1e-3,
+        hc_temperature: float = 0.1
+    ) -> None:
+        super().__init__()
+        self.config = {
+            'input_dim': input_dim, 'num_outputs': num_outputs, 'task_type': task_type,
+            'feature_names': list(feature_names) if feature_names is not None else [f"Feature_{i}" for i in range(input_dim)],
+            'num_facts': num_facts, 'rules_per_layer': tuple(rules_per_layer),
+            'use_calibrators': use_calibrators, 'rule_selection_method': rule_selection_method,
+            'use_prototypes': bool(use_prototypes and task_type == "classification"),
+            'l0_lambda': l0_lambda, 'hc_temperature': hc_temperature
+        }
+
+        if self.config['use_calibrators']:
+            self.calibrators = nn.ModuleList([PiecewiseLinearCalibrator() for _ in range(input_dim)])
+        else:
+            self.calibrators = None
+
+        self.fact = BetaFactLayer(input_dim, num_facts)
+
+        blocks: List[nn.Module] = []
+        cur = num_facts
+        for r in rules_per_layer:
+            if rule_selection_method == 'fixed':
+                blocks.append(SimpleNousBlock(cur, r))
+            elif rule_selection_method == 'softmax':
+                blocks.append(SoftmaxRuleLayer(cur, r))
+            elif rule_selection_method == 'sparse':
+                blocks.append(SparseRuleLayer(cur, r, l0_lambda=l0_lambda, hc_temperature=hc_temperature))
+            else:
+                raise ValueError(f"Unknown rule_selection_method: {rule_selection_method}")
+            cur = r
+        self.blocks = nn.ModuleList(blocks)
+
+        if self.config['use_prototypes']:
+            self.head = ScaledPrototypeLayer(cur, num_prototypes=10, num_classes=num_outputs)
+        else:
+            self.head = nn.Linear(cur, num_outputs)
+
+        if self.config['task_type'] == "regression" and self.config['num_outputs'] != 1:
+            self.config['num_outputs'] = 1
+            if isinstance(self.head, nn.Linear):
+                self.head = nn.Linear(cur, 1)
+
+    def forward(self, x: torch.Tensor, return_internals: bool = False):
+        internals: Dict[str, torch.Tensor] = {}
+        if self.calibrators is not None:
+            x = torch.stack([calib(x[:, i]) for i, calib in enumerate(self.calibrators)], dim=1)
+
+        h_facts = self.fact(x)
+        h = h_facts
+        if return_internals:
+            internals['facts'] = h_facts.detach()
+
+        for i, blk in enumerate(self.blocks):
+            if return_internals:
+                h, details = blk(h, return_details=True)
+                internals[f'block_{i}'] = details
+            else:
+                h = blk(h)
+
+        logits = self.head(h)
+        if self.config['task_type'] == "regression":
+            logits = logits.squeeze(-1)
+
+        if return_internals:
+            return logits, internals
+        return logits
+
+    @torch.no_grad()
+    def forward_explain(
+        self,
+        x: Union[np.ndarray, torch.Tensor],
+        drop_rule_spec: Optional[Tuple[int, int]] = None,
+        restrict_masks: Optional[List[torch.Tensor]] = None,
+        apply_pruning: bool = False,
+        pruning_threshold: float = 0.0,
+        device: Optional[torch.device] = None,
+        explain_disable_norm: bool = False,
+        explain_exclude_proj: bool = False
+    ):
+        """
+        Honest forward for explanations with interventions/gating recompute.
+
+        Returns:
+          - classification: (probas, logits, internals)
+          - regression: (pred, pred, internals)
+        """
+        self.eval()
+        device = device or next(self.parameters()).device
+
+        if isinstance(x, np.ndarray):
+            x = torch.tensor(x, dtype=torch.float32, device=device).unsqueeze(0)
+        elif isinstance(x, torch.Tensor) and x.dim() == 1:
+            x = x.unsqueeze(0).to(device)
+        else:
+            x = x.to(device)
+
+        if self.calibrators is not None:
+            x_cal = torch.stack([calib(x[:, i]) for i, calib in enumerate(self.calibrators)], dim=1)
+        else:
+            x_cal = x
+
+        h = self.fact(x_cal)
+        internals: Dict[str, torch.Tensor] = {'facts': h.detach()}
+
+        for i, blk in enumerate(self.blocks):
+            drop_idx = None
+            if drop_rule_spec is not None and drop_rule_spec[0] == i:
+                drop_idx = int(drop_rule_spec[1])
+            restrict = None
+            if restrict_masks is not None and i < len(restrict_masks) and restrict_masks[i] is not None:
+                restrict = restrict_masks[i].to(device)
+            prune = pruning_threshold if apply_pruning else None
+
+            h, details = blk(
+                h,
+                return_details=True,
+                drop_rule_idx=drop_idx,
+                restrict_mask=restrict,
+                prune_below=prune,
+                explain_disable_norm=explain_disable_norm,
+                explain_exclude_proj=explain_exclude_proj
+            )
+            internals[f'block_{i}'] = details
+
+        logits = self.head(h)
+        if self.config['task_type'] == "classification":
+            probas = F.softmax(logits, dim=-1)
+            return probas.squeeze(0).cpu().numpy(), logits.squeeze(0).cpu().numpy(), internals
+        else:
+            pred = logits.squeeze(-1).squeeze(0).cpu().numpy()
+            return np.array([pred]), np.array([pred]), internals
+
+    def compute_total_l0_loss(self) -> torch.Tensor:
+        if self.config['rule_selection_method'] != 'sparse':
+            return torch.tensor(0.0, device=next(self.parameters()).device)
+        total_l0 = 0.0
+        for blk in self.blocks:
+            if hasattr(blk, 'compute_l0_loss'):
+                total_l0 += blk.compute_l0_loss()
+        return total_l0
+
+    def model_summary(self) -> Dict[str, object]:
+        summary = {
+            "Task": self.config['task_type'],
+            "Rule Selection": self.config['rule_selection_method'],
+            "Use Calibrators": self.config['use_calibrators'],
+            "Use Prototypes": self.config['use_prototypes'],
+            "Num Facts": self.config['num_facts'],
+            "Rules per Layer": self.config['rules_per_layer'],
+            "Total Parameters": sum(p.numel() for p in self.parameters())
+        }
+        if self.config['rule_selection_method'] == 'sparse':
+            summary["L0 Lambda"] = self.config['l0_lambda']
+        return summary
+
+    @torch.no_grad()
+    def encode(
+        self,
+        X: Union[np.ndarray, torch.Tensor, Sequence[Sequence[float]]],
+        device: Optional[torch.device] = None,
+        batch_size: int = 2048,
+        explain_disable_norm: bool = False,
+        explain_exclude_proj: bool = False
+    ) -> torch.Tensor:
+        """
+        Return H [N, D_last] — representations at the head input.
+
+        Flags allow clean representations (without LayerNorm / without residual projection).
+        """
+        self.eval()
+        device = device or next(self.parameters()).device
+
+        if isinstance(X, np.ndarray):
+            X_tensor = torch.tensor(X, dtype=torch.float32)
+        elif isinstance(X, torch.Tensor):
+            X_tensor = X.detach().cpu().float()
+        else:
+            X_tensor = torch.tensor(np.asarray(X), dtype=torch.float32)
+
+        H_list = []
+        for i in range(0, len(X_tensor), batch_size):
+            xb = X_tensor[i:i+batch_size].to(device)
+
+            if self.calibrators is not None:
+                xb_cal = torch.stack([calib(xb[:, j]) for j, calib in enumerate(self.calibrators)], dim=1)
+            else:
+                xb_cal = xb
+
+            h = self.fact(xb_cal)
+            for blk in self.blocks:
+                h, _ = blk(h, return_details=True,
+                           explain_disable_norm=explain_disable_norm,
+                           explain_exclude_proj=explain_exclude_proj)
+            H_list.append(h.detach().cpu())
+        H = torch.cat(H_list, dim=0)
+        return H

nous/prototypes.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict
+
+class ScaledPrototypeLayer(nn.Module):
+    """
+    Prototype-based head for classification with temperature scaling on distances.
+    Uses L2-normalized features/prototypes and exp(-tau * distance).
+    """
+    def __init__(self, input_dim: int, num_prototypes: int, num_classes: int) -> None:
+        super().__init__()
+        self.num_prototypes = num_prototypes
+        self.num_classes = num_classes
+        self.temperature = nn.Parameter(torch.tensor(1.0))
+        self.prototypes = nn.Parameter(torch.randn(num_prototypes, input_dim) * 0.1)
+        self.prototype_class = nn.Parameter(torch.randn(num_prototypes, num_classes))
+
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        h_norm = F.normalize(h, p=2, dim=1)
+        prototypes_norm = F.normalize(self.prototypes, p=2, dim=1)
+        distances = torch.cdist(h_norm, prototypes_norm)
+        activations = torch.exp(-F.softplus(self.temperature) * distances)
+        class_logits = torch.matmul(activations, self.prototype_class)
+        return class_logits
+
+    @torch.no_grad()
+    def get_params(self) -> Dict[str, torch.Tensor]:
+        P = self.prototypes.detach()
+        Pn = F.normalize(P, p=2, dim=1)
+        W = self.prototype_class.detach()
+        Wsm = F.softmax(W, dim=1)
+        tau = F.softplus(self.temperature).detach()
+        return dict(prototypes=P, prototypes_norm=Pn, class_weights=W, class_probs=Wsm, temperature=tau)
+
+    @torch.no_grad()
+    def compute_dist_act(self, h: torch.Tensor):
+        h_norm = F.normalize(h, p=2, dim=1)
+        Pn = F.normalize(self.prototypes, p=2, dim=1)
+        d = torch.cdist(h_norm, Pn)
+        act = torch.exp(-F.softplus(self.temperature) * d)
+        return d, act

nous/rules/__init__.py

@@ -0,0 +1,11 @@
+from .fixed import FixedPairRuleLayer
+from .softmax import SoftmaxRuleLayer
+from .sparse import SparseRuleLayer
+from .blocks import SimpleNousBlock
+
+__all__ = [
+    "FixedPairRuleLayer",
+    "SoftmaxRuleLayer",
+    "SparseRuleLayer",
+    "SimpleNousBlock",
+]

nous/rules/blocks.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+import torch
+import torch.nn as nn
+from .fixed import FixedPairRuleLayer
+
+class SimpleNousBlock(nn.Module):
+    """
+    Residual block over FixedPairRuleLayer with honest interventions and explain flags.
+
+    Options:
+      - drop_rule_idx: zero a single rule pre-gating
+      - restrict_mask: 0/1 mask over rules (frozen active set)
+      - prune_below: threshold prune post-gating
+      - explain_disable_norm / explain_exclude_proj for clean contributions
+    """
+    def __init__(self, in_dim: int, n_rules: int) -> None:
+        super().__init__()
+        self.rule = FixedPairRuleLayer(in_dim, n_rules)
+        self.proj = nn.Identity() if in_dim == n_rules else nn.Linear(in_dim, n_rules, bias=False)
+        self.norm = nn.LayerNorm(n_rules)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        return_details: bool = False,
+        drop_rule_idx: int | None = None,
+        restrict_mask: torch.Tensor | None = None,
+        prune_below: float | None = None,
+        explain_disable_norm: bool = False,
+        explain_exclude_proj: bool = False
+    ):
+        rule_activations = self.rule(x)
+
+        if restrict_mask is not None:
+            rule_activations = rule_activations * restrict_mask
+        if drop_rule_idx is not None:
+            rule_activations[:, drop_rule_idx] = 0.0
+
+        gated_activations = rule_activations
+        if prune_below is not None:
+            mask = (gated_activations.abs() >= prune_below).float()
+            gated_activations = gated_activations * mask
+
+        proj_contrib = torch.zeros_like(gated_activations)
+        if not isinstance(self.proj, nn.Identity):
+            proj_contrib = self.proj(x)
+
+        pre_sum = (0.0 * proj_contrib) + gated_activations if explain_exclude_proj else (proj_contrib + gated_activations)
+        output = pre_sum if explain_disable_norm else self.norm(pre_sum)
+
+        if return_details:
+            details = {
+                "pre_rule_activations": rule_activations.detach(),
+                "gated_activations": gated_activations.detach(),
+                "gate_mask": torch.ones_like(gated_activations, dtype=torch.float32),
+                "aggregator_weights": None,
+                "selected_indices": torch.arange(gated_activations.shape[1], device=x.device).unsqueeze(0).repeat(x.size(0), 1),
+                "facts_used": self.rule.idx.detach(),
+                "pre_norm_sum": pre_sum.detach(),
+                "proj_contrib": proj_contrib.detach()
+            }
+            return output, details
+        return output

nous/rules/fixed.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+import torch
+import torch.nn as nn
+
+class FixedPairRuleLayer(nn.Module):
+    """
+    Fixed random fact pairs with AND logic. Pair indices are distinct per rule.
+    Output: rule_strength * (f_i * f_j)
+    """
+    def __init__(self, input_dim: int, num_rules: int) -> None:
+        super().__init__()
+        idx1 = torch.randint(0, input_dim, (num_rules,))
+        offset = torch.randint(1, input_dim, (num_rules,))
+        idx2 = (idx1 + offset) % input_dim
+        idx = torch.stack([idx1, idx2], dim=1)
+        self.register_buffer('idx', idx)
+        self.weight = nn.Parameter(torch.ones(num_rules))
+
+    def forward(self, facts: torch.Tensor) -> torch.Tensor:
+        f1, f2 = facts[:, self.idx[:,0]], facts[:, self.idx[:,1]]
+        rule_strength = torch.sigmoid(self.weight)
+        return rule_strength * (f1 * f2)
+
+    @torch.no_grad()
+    def get_rules(self):
+        return self.idx.cpu().numpy(), torch.sigmoid(self.weight).cpu().numpy()

nous/rules/softmax.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class SoftmaxRuleLayer(nn.Module):
+    """
+    Learnable fact connections via softmax gating with per-rule top-k and aggregator mixing.
+    Aggregators: AND, OR, k-of-n. Honest interventions supported.
+    """
+    def __init__(self, input_dim: int, num_rules: int, top_k_facts: int = 2, top_k_rules: int = 8) -> None:
+        super().__init__()
+        self.input_dim = input_dim
+        self.num_rules = num_rules
+        self.top_k_facts = top_k_facts
+        self.top_k_rules = top_k_rules
+
+        self.fact_logits = nn.Parameter(torch.randn(num_rules, input_dim) * 0.01)
+        self.num_aggregators = 3
+        self.aggregator_logits = nn.Parameter(torch.zeros(num_rules, self.num_aggregators))
+        self.rule_strength_raw = nn.Parameter(torch.zeros(num_rules))
+        self.proj = nn.Linear(input_dim, num_rules, bias=False) if input_dim != num_rules else nn.Identity()
+        self.norm = nn.LayerNorm(num_rules)
+
+    def forward(
+        self,
+        facts: torch.Tensor,
+        return_details: bool = False,
+        drop_rule_idx: int | None = None,
+        restrict_mask: torch.Tensor | None = None,
+        prune_below: float | None = None,
+        explain_disable_norm: bool = False,
+        explain_exclude_proj: bool = False
+    ):
+        fact_logits_soft = F.softmax(self.fact_logits, dim=1)
+        _, topk_fact_idx = torch.topk(fact_logits_soft, k=min(self.top_k_facts, self.input_dim), dim=1)
+        mask = torch.zeros_like(fact_logits_soft)
+        mask.scatter_(1, topk_fact_idx, 1.0)
+
+        facts_expanded = facts.unsqueeze(1)
+        mask_expanded = mask.unsqueeze(0)
+        selected_facts = facts_expanded * mask_expanded
+
+        and_agg = torch.prod(selected_facts + (1 - mask_expanded), dim=2)
+        or_agg = 1 - torch.prod((1 - selected_facts) * mask_expanded + (1 - mask_expanded), dim=2)
+        k_of_n_agg = torch.sum(selected_facts, dim=2) / (mask_expanded.sum(dim=2) + 1e-8)
+
+        agg_weights = F.softmax(self.aggregator_logits, dim=1)
+        aggregators = torch.stack([and_agg, or_agg, k_of_n_agg], dim=2)
+        mixed_agg = torch.sum(aggregators * agg_weights.unsqueeze(0), dim=2)
+
+        rule_strength = torch.sigmoid(self.rule_strength_raw)
+        rule_activations = mixed_agg * rule_strength.unsqueeze(0)
+
+        pre_for_topk = rule_activations.clone()
+        if restrict_mask is not None:
+            pre_for_topk = pre_for_topk + (restrict_mask - 1) * 1e9
+        if drop_rule_idx is not None:
+            pre_for_topk[:, drop_rule_idx] = -1e9
+
+        k = min(self.top_k_rules, self.num_rules)
+        _, topk_rule_idx = torch.topk(pre_for_topk, k=k, dim=1)
+        gate_mask = torch.zeros_like(rule_activations)
+        gate_mask.scatter_(1, topk_rule_idx, 1.0)
+
+        if restrict_mask is not None:
+            gate_mask = gate_mask * restrict_mask.unsqueeze(0).to(gate_mask.dtype)
+        if drop_rule_idx is not None:
+            gate_mask[:, drop_rule_idx] = 0.0
+
+        gated_activations = rule_activations * gate_mask
+        if prune_below is not None:
+            keep = (gated_activations.abs() >= prune_below).float()
+            gated_activations = gated_activations * keep
+            gate_mask = gate_mask * keep
+
+        proj_contrib = self.proj(facts) if not isinstance(self.proj, nn.Identity) else facts
+        pre_sum = gated_activations if explain_exclude_proj else (proj_contrib + gated_activations)
+        output = pre_sum if explain_disable_norm else self.norm(pre_sum)
+
+        if return_details:
+            details = {
+                "pre_rule_activations": rule_activations.detach(),
+                "gated_activations": gated_activations.detach(),
+                "gate_mask": gate_mask.detach(),
+                "aggregator_weights": agg_weights.detach(),
+                "selected_indices": topk_rule_idx.detach(),
+                "facts_used": topk_fact_idx.detach(),
+                "pre_norm_sum": pre_sum.detach(),
+                "proj_contrib": proj_contrib.detach()
+            }
+            return output, details
+        return output

nous/rules/sparse.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class HardConcretePerConnection(nn.Module):
+    """
+    Per-connection Hard-Concrete gate with safer sampling (clamped u and log_sigma).
+    """
+    def __init__(self, shape, init_mean=0.0, init_std=1.0, temperature=0.1, stretch_limits=(-0.1, 1.1)):
+        super().__init__()
+        self.shape = shape
+        self.beta = nn.Parameter(torch.ones(shape) * init_mean)
+        self.log_sigma = nn.Parameter(torch.ones(shape) * init_std)
+        self.temperature = temperature
+        self.l, self.r = stretch_limits
+
+    def forward(self, training: bool | None = None) -> torch.Tensor:
+        if training is None:
+            training = self.training
+        if training:
+            u = torch.rand(self.shape, device=self.beta.device).clamp(1e-6, 1 - 1e-6)
+            sigma = torch.exp(torch.clamp(self.log_sigma, -5.0, 5.0))
+            pre = (torch.log(u) - torch.log(1 - u) + self.beta) / (sigma + 1e-8)
+            s = torch.sigmoid(pre)
+            s = s * (self.r - self.l) + self.l
+            left  = torch.sigmoid((s - self.l) / self.temperature)
+            right = torch.sigmoid((s - self.r) / self.temperature)
+            mask = (left - right).clamp(0.0, 1.0)
+        else:
+            mask = (torch.sigmoid(self.beta) > 0.5).float()
+        return mask
+
+    def get_proba(self) -> torch.Tensor:
+        l_tensor = torch.tensor(self.l, device=self.beta.device)
+        r_tensor = torch.tensor(self.r, device=self.beta.device)
+        return torch.sigmoid(torch.clamp(self.beta, -10.0, 10.0) - self.temperature * torch.log(-l_tensor / r_tensor + 1e-8))
+
+    def l0_penalty(self) -> torch.Tensor:
+        return self.get_proba().sum()
+
+
+class SparseRuleLayer(nn.Module):
+    """
+    Hard-Concrete sparse connections + aggregator mixing + honest interventions.
+    Aggregators: AND, OR, k-of-n, NOT.
+    """
+    def __init__(self, input_dim: int, num_rules: int, top_k_facts: int = 2, top_k_rules: int = 8, l0_lambda: float = 1e-3, hc_temperature: float = 0.1):
+        super().__init__()
+        self.input_dim = input_dim
+        self.num_rules = num_rules
+        self.top_k_facts = top_k_facts
+        self.top_k_rules = top_k_rules
+        self.l0_lambda = l0_lambda
+
+        self.hard_concrete = HardConcretePerConnection(
+            shape=(num_rules, input_dim),
+            init_mean=0.0,
+            init_std=1.0,
+            temperature=hc_temperature,
+            stretch_limits=(-0.1, 1.1)
+        )
+
+        self.num_aggregators = 4
+        self.aggregator_logits = nn.Parameter(torch.zeros(num_rules, self.num_aggregators))
+        self.rule_strength_raw = nn.Parameter(torch.zeros(num_rules))
+        self.proj = nn.Linear(input_dim, num_rules, bias=False) if input_dim != num_rules else nn.Identity()
+        self.norm = nn.LayerNorm(num_rules)
+
+    def forward(
+        self,
+        facts: torch.Tensor,
+        return_details: bool = False,
+        drop_rule_idx: int | None = None,
+        restrict_mask: torch.Tensor | None = None,
+        prune_below: float | None = None,
+        explain_disable_norm: bool = False,
+        explain_exclude_proj: bool = False
+    ):
+        mask = self.hard_concrete(training=self.training)  # [R, F]
+        facts_expanded = facts.unsqueeze(1)
+        mask_expanded = mask.unsqueeze(0)
+        selected_facts = facts_expanded * mask_expanded
+
+        and_agg = torch.prod(selected_facts + (1 - mask_expanded), dim=2)
+        or_agg = 1 - torch.prod((1 - selected_facts) * mask_expanded + (1 - mask_expanded), dim=2)
+        selected_count = torch.sum(mask_expanded, dim=2) + 1e-8
+        k_of_n_agg = torch.sum(selected_facts, dim=2) / selected_count
+        not_agg = 1 - k_of_n_agg
+
+        agg_weights = F.softmax(self.aggregator_logits, dim=1)  # [R, 4]
+        aggregators = torch.stack([and_agg, or_agg, k_of_n_agg, not_agg], dim=2)
+        mixed_agg = torch.sum(aggregators * agg_weights.unsqueeze(0), dim=2)  # [B, R]
+
+        rule_strength = torch.sigmoid(self.rule_strength_raw)  # [R]
+        rule_activations = mixed_agg * rule_strength.unsqueeze(0)  # [B, R]
+
+        pre_for_topk = rule_activations.clone()
+        if restrict_mask is not None:
+            pre_for_topk = pre_for_topk + (restrict_mask - 1) * 1e9
+        if drop_rule_idx is not None:
+            pre_for_topk[:, drop_rule_idx] = -1e9
+
+        k = min(self.top_k_rules, self.num_rules)
+        _, topk_rule_idx = torch.topk(pre_for_topk, k=k, dim=1)
+        gate_mask = torch.zeros_like(rule_activations)
+        gate_mask.scatter_(1, topk_rule_idx, 1.0)
+
+        if restrict_mask is not None:
+            gate_mask = gate_mask * restrict_mask.unsqueeze(0).to(gate_mask.dtype)
+        if drop_rule_idx is not None:
+            gate_mask[:, drop_rule_idx] = 0.0
+
+        gated_activations = rule_activations * gate_mask
+        if prune_below is not None:
+            keep = (gated_activations.abs() >= prune_below).float()
+            gated_activations = gated_activations * keep
+            gate_mask = gate_mask * keep
+
+        proj_contrib = self.proj(facts) if not isinstance(self.proj, nn.Identity) else facts
+        pre_sum = gated_activations if explain_exclude_proj else (proj_contrib + gated_activations)
+        output = pre_sum if explain_disable_norm else self.norm(pre_sum)
+
+        if return_details:
+            with torch.no_grad():
+                _, topk_fact_idx = torch.topk(mask, k=min(self.top_k_facts, self.input_dim), dim=1)
+            details = {
+                "pre_rule_activations": rule_activations.detach(),
+                "gated_activations": gated_activations.detach(),
+                "gate_mask": gate_mask.detach(),
+                "aggregator_weights": agg_weights.detach(),
+                "selected_indices": topk_rule_idx.detach(),
+                "facts_used": topk_fact_idx.detach(),
+                "pre_norm_sum": pre_sum.detach(),
+                "proj_contrib": proj_contrib.detach()
+            }
+            return output, details
+        return output
+
+    def compute_l0_loss(self) -> torch.Tensor:
+        num_conn = float(self.num_rules * self.input_dim)
+        return (self.l0_lambda * self.hard_concrete.l0_penalty()) / max(1.0, num_conn)

nous/training/__init__.py

@@ -0,0 +1,5 @@
+from .train import train_model
+from .evaluation import evaluate_classification, evaluate_regression
+from .schedulers import make_sparse_regression_hook
+
+__all__ = ["train_model", "evaluate_classification", "evaluate_regression", "make_sparse_regression_hook"]

nous/training/evaluation.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from sklearn.metrics import (
+    roc_auc_score, accuracy_score, mean_absolute_error,
+    mean_squared_error, r2_score
+)
+
+def evaluate_classification(model, test_loader: DataLoader, device, class_names=None):
+    """
+    Evaluate classification metrics on a dataloader: accuracy and one-vs-rest AUC.
+    """
+    model.eval()
+    all_probs, all_labels = [], []
+    with torch.no_grad():
+        for X_batch, y_batch in test_loader:
+            X_batch = X_batch.to(device)
+            probs = F.softmax(model(X_batch), dim=1).cpu().numpy()
+            all_probs.extend(probs)
+            all_labels.extend(y_batch.numpy())
+    all_probs = np.array(all_probs)
+    all_preds = np.argmax(all_probs, axis=1)
+    acc = accuracy_score(all_labels, all_preds)
+    try:
+        auc = roc_auc_score(all_labels, all_probs, multi_class='ovr')
+    except ValueError:
+        auc = 0.5
+    return acc, auc, all_probs, all_labels
+
+def evaluate_regression(model, test_loader: DataLoader, device, y_scaler=None):
+    """
+    Evaluate regression: RMSE, MAE, R^2. Returns unscaled predictions if y_scaler is passed.
+    """
+    model.eval()
+    preds, labels = [], []
+    with torch.no_grad():
+        for X_batch, y_batch in test_loader:
+            X_batch = X_batch.to(device)
+            pred = model(X_batch).cpu().numpy().ravel()
+            preds.extend(pred)
+            labels.extend(y_batch.numpy().ravel())
+
+    preds = np.array(preds)
+    labels = np.array(labels)
+
+    if y_scaler is not None:
+        preds_u = y_scaler.inverse_transform(preds.reshape(-1,1)).ravel()
+        labels_u = y_scaler.inverse_transform(labels.reshape(-1,1)).ravel()
+    else:
+        preds_u, labels_u = preds, labels
+
+    rmse = np.sqrt(mean_squared_error(labels_u, preds_u))
+    mae  = mean_absolute_error(labels_u, preds_u)
+    r2   = r2_score(labels_u, preds_u)
+    return rmse, mae, r2, preds_u, labels_u