l0-python 0.4.3__py3-none-any.whl → 0.6.0__py3-none-any.whl

l0/__init__.py

@@ -5,7 +5,7 @@ A PyTorch implementation of L0 regularization for neural network sparsification
 and intelligent sampling, based on Louizos, Welling, & Kingma (2017).
 """
 
-__version__ = "0.1.0"
+__version__ = "0.5.0"
 
 from .distributions import HardConcrete
 from .gates import FeatureGate, HybridGate, L0Gate, SampleGate

l0/calibration.py

@@ -60,6 +60,7 @@ class SparseCalibrationWeights(nn.Module):
         log_weight_jitter_sd: float = 0.0,
         log_alpha_jitter_sd: float = 0.01,
         device: str | torch.device = "cpu",
+        use_gates: bool = True,
     ):
         super().__init__()
         self.n_features = n_features
@@ -69,6 +70,7 @@ class SparseCalibrationWeights(nn.Module):
         self.log_weight_jitter_sd = log_weight_jitter_sd
         self.log_alpha_jitter_sd = log_alpha_jitter_sd
         self.device = torch.device(device)
+        self.use_gates = use_gates
 
         # Initialize weights (on original scale)
         if init_weights is None:
@@ -193,6 +195,10 @@ class SparseCalibrationWeights(nn.Module):
         torch.Tensor
             Positive calibration weights with L0 sparsity applied
         """
+        if not self.use_gates:
+            # No gates - just return positive weights
+            return torch.exp(self.log_weight)
+
         # Sample or get deterministic gates
         if deterministic:
             gates = self.get_deterministic_gates()
@@ -304,6 +310,8 @@ class SparseCalibrationWeights(nn.Module):
         verbose: bool = False,
         verbose_freq: int = 100,
         target_groups: np.ndarray | None = None,
+        normalize_groups: bool = True,
+        group_multipliers: dict[int, float] | None = None,
     ) -> "SparseCalibrationWeights":
         """
         Fit calibration weights using gradient descent.
@@ -332,6 +340,12 @@ class SparseCalibrationWeights(nn.Module):
             Array of group IDs for each target. Targets in the same group
             will be averaged together so each group contributes equally to loss.
             If None, all targets are treated independently.
+        normalize_groups : bool
+            Whether to normalize within groups so each group contributes
+            equally to the loss. Default True (backward-compatible).
+        group_multipliers : dict[int, float], optional
+            Per-group loss scaling factors. Applied after normalization
+            (if enabled). Requires target_groups to be set.
 
         Returns
         -------
@@ -344,27 +358,53 @@ class SparseCalibrationWeights(nn.Module):
         # Convert M to torch sparse (will be cached)
         M_torch = self._convert_sparse_to_torch(M)
 
+        # Validate group_multipliers
+        if group_multipliers is not None and target_groups is None:
+            raise ValueError(
+                "group_multipliers requires target_groups to be set"
+            )
+
         # Compute group weights for loss averaging
         if target_groups is not None:
-            # Convert to tensor
             target_groups = torch.tensor(
                 target_groups, dtype=torch.long, device=self.device
             )
-
-            # Calculate group weights: 1 / group_size for each target
             unique_groups = torch.unique(target_groups)
-            group_weights = torch.zeros_like(y)
-
-            for group_id in unique_groups:
-                group_mask = target_groups == group_id
-                group_size = group_mask.sum().item()
-                # Each target in the group gets weight 1/group_size
-                # so the group's total contribution is 1
-                group_weights[group_mask] = 1.0 / group_size
+
+            if normalize_groups:
+                group_weights = torch.zeros_like(y)
+                for group_id in unique_groups:
+                    group_mask = target_groups == group_id
+                    group_size = group_mask.sum().item()
+                    group_weights[group_mask] = 1.0 / group_size
+            else:
+                group_weights = torch.ones_like(y)
+
+            if group_multipliers is not None:
+                for gid, mult in group_multipliers.items():
+                    group_mask = target_groups == gid
+                    if not group_mask.any():
+                        raise ValueError(
+                            f"group_multipliers key {gid} not found "
+                            f"in target_groups"
+                        )
+                    group_weights[group_mask] *= mult
         else:
-            # No grouping - all targets weighted equally
             group_weights = torch.ones_like(y)
 
+        if verbose:
+            if target_groups is not None:
+                n_groups = len(torch.unique(target_groups))
+                parts = [f"{n_groups} groups"]
+                parts.append(
+                    f"normalize={'on' if normalize_groups else 'off'}"
+                )
+                if group_multipliers is not None:
+                    parts.append(f"multipliers={group_multipliers}")
+                print(f"Groups: {', '.join(parts)}")
+            else:
+                print("Groups: none (uniform weights)")
+
         # Add jitter to weights to break symmetry (if jitter_sd > 0)
         if self.log_weight_jitter_sd > 0:
             with torch.no_grad():
@@ -453,6 +493,46 @@ class SparseCalibrationWeights(nn.Module):
                         1 - active_info["count"] / self.n_features
                     )
 
+                    # Weight distribution diagnostic
+                    if active_weights.numel() > 0:
+                        w_tiny = (active_weights < 0.01).sum().item()
+                        w_small = (
+                            ((active_weights >= 0.01) & (active_weights < 0.1))
+                            .sum()
+                            .item()
+                        )
+                        w_med = (
+                            ((active_weights >= 0.1) & (active_weights < 1.0))
+                            .sum()
+                            .item()
+                        )
+                        w_normal = (
+                            ((active_weights >= 1.0) & (active_weights < 10.0))
+                            .sum()
+                            .item()
+                        )
+                        w_large = (
+                            (
+                                (active_weights >= 10.0)
+                                & (active_weights < 1000.0)
+                            )
+                            .sum()
+                            .item()
+                        )
+                        w_huge = (active_weights >= 1000.0).sum().item()
+
+                        total_active = active_weights.numel()
+                        weight_dist = (
+                            f"[<0.01: {100*w_tiny/total_active:.1f}%, "
+                            f"0.01-0.1: {100*w_small/total_active:.1f}%, "
+                            f"0.1-1: {100*w_med/total_active:.1f}%, "
+                            f"1-10: {100*w_normal/total_active:.1f}%, "
+                            f"10-1000: {100*w_large/total_active:.1f}%, "
+                            f">1000: {100*w_huge/total_active:.1f}%]"
+                        )
+                    else:
+                        weight_dist = "[no active weights]"
+
                     # Calculate components of the actual loss being minimized
                     actual_data_loss = data_loss.item()
                     actual_l0_loss = l0_loss.item()
@@ -464,7 +544,8 @@ class SparseCalibrationWeights(nn.Module):
                             f"mean_group_mare={mean_group_mare:.4%}, "
                             f"max_error={max_rel_err:.1%}, "
                             f"total_loss={actual_total_loss:.3f}, "
-                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
+                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)\n"
+                            f"         Weight dist: {weight_dist}"
                         )
                     else:
                         print(
@@ -472,7 +553,8 @@ class SparseCalibrationWeights(nn.Module):
                             f"mean_error={mean_rel_err:.4%}, "
                             f"max_error={max_rel_err:.1%}, "
                             f"total_loss={actual_total_loss:.3f}, "
-                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)"
+                            f"active={active_info['count']:4d}/{self.n_features} ({sparsity_pct:.1f}% sparse)\n"
+                            f"         Weight dist: {weight_dist}"
                         )
 
         return self

{l0_python-0.4.3.dist-info → l0_python-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.4.3
+Version: 0.6.0
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License

l0_python-0.6.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+l0/__init__.py,sha256=ImRgZfVg-RJhLoqNaLc8ysjddZ7CYGHJPJbRUyfSoHY,1192
+l0/calibration.py,sha256=xOSkuqU72rfFV8OJi-OntYGcPwqabDGrQ12VHBZqNK0,21514
+l0/distributions.py,sha256=GvVoDasdxoDlPPdzt33lFzkuJgQsfkEn9e-abxru3nM,6325
+l0/gates.py,sha256=54OORMg9qSsdbIyzYIL9t9CCHWsKbumSKl_d4-YFuNU,11443
+l0/layers.py,sha256=j_DPEFdfUc37ojtoK2ufGzVYzZ_DwusqTcJXCVxZS68,13785
+l0/penalties.py,sha256=okqXgX2eGUmg5prcDclicTQALPPPF9FA5eQP09V4GTE,9569
+l0/sparse.py,sha256=UVpkrrJBnQ9iSNDoWK5F_sr43tQareWXzYHyMq_cKss,9726
+l0_python-0.6.0.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
+l0_python-0.6.0.dist-info/METADATA,sha256=81inHMQ72BeOFNDDkYvZSzZ9C75e23rSprrBjFPxAFw,7095
+l0_python-0.6.0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+l0_python-0.6.0.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
+l0_python-0.6.0.dist-info/RECORD,,

{l0_python-0.4.3.dist-info → l0_python-0.6.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (82.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

l0_python-0.4.3.dist-info/RECORD

@@ -1,12 +0,0 @@
-l0/__init__.py,sha256=ezGVLGbdU_jAElAJSdLKjSWCnZRILAZi8qMBPjcHwQc,1192
-l0/calibration.py,sha256=3VM-28M1yp-DM6AlHX7sqZCaE7OJevpY9Ey-syjSPik,17838
-l0/distributions.py,sha256=GvVoDasdxoDlPPdzt33lFzkuJgQsfkEn9e-abxru3nM,6325
-l0/gates.py,sha256=54OORMg9qSsdbIyzYIL9t9CCHWsKbumSKl_d4-YFuNU,11443
-l0/layers.py,sha256=j_DPEFdfUc37ojtoK2ufGzVYzZ_DwusqTcJXCVxZS68,13785
-l0/penalties.py,sha256=okqXgX2eGUmg5prcDclicTQALPPPF9FA5eQP09V4GTE,9569
-l0/sparse.py,sha256=UVpkrrJBnQ9iSNDoWK5F_sr43tQareWXzYHyMq_cKss,9726
-l0_python-0.4.3.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
-l0_python-0.4.3.dist-info/METADATA,sha256=gN2ufd7zkVDielkTS2iGbUeKkf9bIMTiGZtvCkpc960,7095
-l0_python-0.4.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-l0_python-0.4.3.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
-l0_python-0.4.3.dist-info/RECORD,,