l0-python 0.6.1__py3-none-any.whl → 0.6.2__py3-none-any.whl

l0/__init__.py

@@ -7,6 +7,7 @@ and intelligent sampling, based on Louizos, Welling, & Kingma (2017).
 
 __version__ = "0.5.0"
 
+from .calibration import SparseCalibrationWeights
 from .distributions import HardConcrete
 from .gates import FeatureGate, HybridGate, L0Gate, SampleGate
 from .layers import (
@@ -39,6 +40,8 @@ __all__ = [
     "prune_model",
     # Sparse
     "SparseL0Linear",
+    # Calibration
+    "SparseCalibrationWeights",
     # Gates
     "L0Gate",
     "SampleGate",

l0/calibration.py

@@ -47,6 +47,12 @@ class SparseCalibrationWeights(nn.Module):
         Set to 0 to disable jitter. Default is 0.01 (following Louizos et al.).
     device : str or torch.device
         Device to run computations on ('cpu' or 'cuda')
+    seed : int, optional
+        Seed for the RNG used by the `log_alpha` init jitter and (inside
+        ``fit``) the `log_weight` jitter. When set, two models with the same
+        inputs produce byte-identical initial `log_alpha` and jitter without
+        the caller having to manage PyTorch's global RNG. ``None`` (default)
+        preserves legacy behaviour of using the global RNG.
     """
 
     def __init__(
@@ -61,6 +67,7 @@ class SparseCalibrationWeights(nn.Module):
         log_alpha_jitter_sd: float = 0.01,
         device: str | torch.device = "cpu",
         use_gates: bool = True,
+        seed: int | None = None,
     ):
         super().__init__()
         self.n_features = n_features
@@ -71,6 +78,17 @@ class SparseCalibrationWeights(nn.Module):
         self.log_alpha_jitter_sd = log_alpha_jitter_sd
         self.device = torch.device(device)
         self.use_gates = use_gates
+        self.seed = seed
+
+        # Local RNG (only used when `seed` is provided). Kept on the module so
+        # `fit`'s `log_weight` jitter uses the same deterministic stream.
+        if seed is not None:
+            self._generator: torch.Generator | None = torch.Generator(
+                device=self.device
+            )
+            self._generator.manual_seed(int(seed))
+        else:
+            self._generator = None
 
         # Initialize weights (on original scale)
         if init_weights is None:
@@ -121,7 +139,7 @@ class SparseCalibrationWeights(nn.Module):
         # Add jitter to break symmetry (if specified)
         if self.log_alpha_jitter_sd > 0:
             jitter = (
-                torch.randn(n_features, device=self.device)
+                torch.randn(n_features, generator=self._generator, device=self.device)
                 * self.log_alpha_jitter_sd
             )
             self.log_alpha = nn.Parameter(mu + jitter)
@@ -132,9 +150,7 @@ class SparseCalibrationWeights(nn.Module):
         self._cached_M_torch: torch.sparse.Tensor | None = None
         self._cached_M_shape: tuple[int, int] | None = None
 
-    def _convert_sparse_to_torch(
-        self, M_sparse: sp.spmatrix
-    ) -> torch.sparse.Tensor:
+    def _convert_sparse_to_torch(self, M_sparse: sp.spmatrix) -> torch.sparse.Tensor:
         """
         Convert scipy sparse matrix to torch sparse tensor.
 
@@ -148,9 +164,7 @@ class SparseCalibrationWeights(nn.Module):
             return self._cached_M_torch
 
         M_coo = M_sparse.tocoo()
-        indices = torch.LongTensor(np.vstack([M_coo.row, M_coo.col])).to(
-            self.device
-        )
+        indices = torch.LongTensor(np.vstack([M_coo.row, M_coo.col])).to(self.device)
         values = torch.FloatTensor(M_coo.data).to(self.device)
         M_torch = torch.sparse_coo_tensor(
             indices,
@@ -166,18 +180,26 @@ class SparseCalibrationWeights(nn.Module):
 
         return M_torch
 
+    # Bounds for `log_alpha` before scaling by `beta`. Values outside
+    # ``[-LOG_ALPHA_BOUND, LOG_ALPHA_BOUND]`` saturate the sigmoid and
+    # vanish gradients; clamping keeps sampling and deterministic gates
+    # well-defined under fp16/bf16.
+    LOG_ALPHA_BOUND: float = 20.0
+
     def _sample_gates(self) -> torch.Tensor:
         """Sample gates using Hard Concrete distribution."""
         eps = 1e-6
-        u = torch.rand_like(self.log_alpha).clamp(eps, 1 - eps)
-        s = (torch.log(u) - torch.log(1 - u) + self.log_alpha) / self.beta
+        log_alpha = self.log_alpha.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        u = torch.rand_like(log_alpha).clamp(eps, 1 - eps)
+        s = (torch.log(u) - torch.log(1 - u) + log_alpha) / self.beta
         s = torch.sigmoid(s)
         s_bar = s * (self.zeta - self.gamma) + self.gamma
         return s_bar.clamp(0, 1)
 
     def get_deterministic_gates(self) -> torch.Tensor:
         """Get deterministic gate values (for inference)."""
-        s = torch.sigmoid(self.log_alpha / self.beta)
+        log_alpha = self.log_alpha.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        s = torch.sigmoid(log_alpha / self.beta)
         s_bar = s * (self.zeta - self.gamma) + self.gamma
         return s_bar.clamp(0, 1)
 
@@ -249,7 +271,8 @@ class SparseCalibrationWeights(nn.Module):
         c = -self.beta * torch.log(
             torch.tensor(-self.gamma / self.zeta, device=self.device)
         )
-        pi = torch.sigmoid(self.log_alpha + c)
+        log_alpha = self.log_alpha.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        pi = torch.sigmoid(log_alpha + c)
         return pi.sum()
 
     def get_l2_penalty(self) -> torch.Tensor:
@@ -360,9 +383,7 @@ class SparseCalibrationWeights(nn.Module):
 
         # Validate group_multipliers
         if group_multipliers is not None and target_groups is None:
-            raise ValueError(
-                "group_multipliers requires target_groups to be set"
-            )
+            raise ValueError("group_multipliers requires target_groups to be set")
 
         # Compute group weights for loss averaging
         if target_groups is not None:
@@ -385,8 +406,7 @@ class SparseCalibrationWeights(nn.Module):
                     group_mask = target_groups == gid
                     if not group_mask.any():
                         raise ValueError(
-                            f"group_multipliers key {gid} not found "
-                            f"in target_groups"
+                            f"group_multipliers key {gid} not found in target_groups"
                         )
                     group_weights[group_mask] *= mult
         else:
@@ -396,9 +416,7 @@ class SparseCalibrationWeights(nn.Module):
             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'}"
-                )
+                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)}")
@@ -408,10 +426,17 @@ class SparseCalibrationWeights(nn.Module):
         # Add jitter to weights to break symmetry (if jitter_sd > 0)
         if self.log_weight_jitter_sd > 0:
             with torch.no_grad():
-                jitter = (
-                    torch.randn_like(self.log_weight)
-                    * self.log_weight_jitter_sd
-                )
+                # `torch.randn_like` can't take a generator kwarg, so draw
+                # explicitly via `torch.randn` when a local RNG is set.
+                if self._generator is not None:
+                    noise = torch.randn(
+                        self.log_weight.shape,
+                        generator=self._generator,
+                        device=self.device,
+                    )
+                else:
+                    noise = torch.randn_like(self.log_weight)
+                jitter = noise * self.log_weight_jitter_sd
                 self.log_weight.data += jitter
 
         # Setup optimizer
@@ -428,9 +453,7 @@ class SparseCalibrationWeights(nn.Module):
                 # Adding 1 to avoid division by zero
                 relative_errors = (y - y_pred) / (y + 1)
                 # Apply group weights and then average
-                weighted_squared_errors = (
-                    relative_errors.pow(2) * group_weights
-                )
+                weighted_squared_errors = relative_errors.pow(2) * group_weights
                 data_loss = (
                     weighted_squared_errors.sum()
                 )  # Sum because weights already normalize
@@ -480,18 +503,14 @@ class SparseCalibrationWeights(nn.Module):
                         group_losses = []
                         for group_id in torch.unique(target_groups):
                             group_mask = target_groups == group_id
-                            group_mean_err = (
-                                rel_errors[group_mask].mean().item()
-                            )
+                            group_mean_err = rel_errors[group_mask].mean().item()
                             group_losses.append(group_mean_err)
                         mean_group_mare = np.mean(group_losses)
                     else:
                         mean_group_mare = mean_rel_err
 
                     # Calculate sparsity percentage
-                    sparsity_pct = 100 * (
-                        1 - active_info["count"] / self.n_features
-                    )
+                    sparsity_pct = 100 * (1 - active_info["count"] / self.n_features)
 
                     # Weight distribution diagnostic
                     if active_weights.numel() > 0:
@@ -512,10 +531,7 @@ class SparseCalibrationWeights(nn.Module):
                             .item()
                         )
                         w_large = (
-                            (
-                                (active_weights >= 10.0)
-                                & (active_weights < 1000.0)
-                            )
+                            ((active_weights >= 10.0) & (active_weights < 1000.0))
                             .sum()
                             .item()
                         )
@@ -523,12 +539,12 @@ class SparseCalibrationWeights(nn.Module):
 
                         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}%]"
+                            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]"
@@ -540,7 +556,7 @@ class SparseCalibrationWeights(nn.Module):
 
                     if target_groups is not None:
                         print(
-                            f"Epoch {epoch+1:4d}: "
+                            f"Epoch {epoch + 1:4d}: "
                             f"mean_group_mare={mean_group_mare:.4%}, "
                             f"max_error={max_rel_err:.1%}, "
                             f"total_loss={actual_total_loss:.3f}, "
@@ -549,7 +565,7 @@ class SparseCalibrationWeights(nn.Module):
                         )
                     else:
                         print(
-                            f"Epoch {epoch+1:4d}: "
+                            f"Epoch {epoch + 1:4d}: "
                             f"mean_error={mean_rel_err:.4%}, "
                             f"max_error={max_rel_err:.1%}, "
                             f"total_loss={actual_total_loss:.3f}, "

l0/distributions.py

@@ -69,9 +69,7 @@ class HardConcrete(nn.Module):
             init_val = math.log(self.init_mean / (1 - self.init_mean))
             self.qz_logits.data.fill_(init_val)
 
-    def forward(
-        self, input_shape: tuple[int, ...] | None = None
-    ) -> torch.Tensor:
+    def forward(self, input_shape: tuple[int, ...] | None = None) -> torch.Tensor:
         """
         Sample or compute gates.
 
@@ -98,6 +96,16 @@ class HardConcrete(nn.Module):
 
         return gates
 
+    # Bounds for `qz_logits` before sigmoid/temperature scaling. Values
+    # outside ``[-LOG_ALPHA_BOUND, LOG_ALPHA_BOUND]`` saturate the sigmoid
+    # and vanish gradients; clamping keeps sampling and deterministic gates
+    # well-defined in fp16/bf16 as well. See Louizos et al. 2017.
+    LOG_ALPHA_BOUND: float = 20.0
+
+    # Epsilon for uniform sampling; fp16-safe (``1e-8`` underflows to zero
+    # at fp16 and ``log(0) = -inf``).
+    _UNIFORM_EPS: float = 1e-6
+
     def _sample_gates(self) -> torch.Tensor:
         """
         Sample gates using the reparameterization trick.
@@ -107,12 +115,19 @@ class HardConcrete(nn.Module):
         torch.Tensor
             Sampled gate values in [0, 1]
         """
-        # Sample uniform noise (avoiding exact 0 and 1 for numerical stability)
-        u = torch.zeros_like(self.qz_logits).uniform_(1e-8, 1.0 - 1e-8)
+        # Clamp logits to a safe range so log-odds stay finite under fp16/bf16
+        # and gradients don't vanish at saturation.
+        logits = self.qz_logits.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+
+        # Sample uniform noise (avoiding exact 0 and 1 for numerical stability).
+        # ``1e-6`` matches calibration.py / sparse.py and is fp16-safe.
+        u = torch.zeros_like(logits).uniform_(
+            self._UNIFORM_EPS, 1.0 - self._UNIFORM_EPS
+        )
 
         # Apply the concrete distribution transformation
         # s = sigmoid((log(u) - log(1-u) + logits) / temperature)
-        s = torch.log(u) - torch.log(1 - u) + self.qz_logits
+        s = torch.log(u) - torch.log(1 - u) + logits
         s = torch.sigmoid(s / self.temperature)
 
         # Stretch and clamp to create hard concrete
@@ -125,13 +140,20 @@ class HardConcrete(nn.Module):
         """
         Compute deterministic gates for evaluation.
 
+        Uses the mean of the Hard Concrete distribution, which applies the
+        temperature scaling to ``qz_logits`` so that ``eval()`` output is
+        consistent with ``_sample_gates`` and ``get_penalty`` /
+        ``get_active_prob`` (see Louizos et al. 2017, Eq. 11).
+
         Returns
         -------
         torch.Tensor
             Deterministic gate values in [0, 1]
         """
-        # Use mean of the distribution
-        probs = torch.sigmoid(self.qz_logits)
+        # Clamp logits so eval output stays well-defined in fp16/bf16.
+        logits = self.qz_logits.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        # Mean of the binary concrete before stretch: sigmoid(logits / beta).
+        probs = torch.sigmoid(logits / self.temperature)
 
         # Apply stretching transformation
         gates = probs * (self.zeta - self.gamma) + self.gamma
@@ -147,10 +169,10 @@ class HardConcrete(nn.Module):
         torch.Tensor
             Expected number of non-zero gates
         """
+        # Clamp logits so sigmoid doesn't saturate and kill gradients.
+        logits = self.qz_logits.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
         # Shift logits to account for hard concrete bounds
-        logits_shifted = self.qz_logits - self.temperature * math.log(
-            -self.gamma / self.zeta
-        )
+        logits_shifted = logits - self.temperature * math.log(-self.gamma / self.zeta)
 
         # Probability that gate is active (non-zero)
         prob_active = torch.sigmoid(logits_shifted)
@@ -166,9 +188,8 @@ class HardConcrete(nn.Module):
         torch.Tensor
             Probability of each gate being non-zero
         """
-        logits_shifted = self.qz_logits - self.temperature * math.log(
-            -self.gamma / self.zeta
-        )
+        logits = self.qz_logits.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        logits_shifted = logits - self.temperature * math.log(-self.gamma / self.zeta)
         return torch.sigmoid(logits_shifted)
 
     def get_sparsity(self) -> float:

l0/gates.py

@@ -130,9 +130,7 @@ class SampleGate(L0Gate):
         self.n_samples = n_samples
         self.target_samples = target_samples
 
-    def select_samples(
-        self, data: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def select_samples(self, data: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Select samples from data.
 
@@ -219,9 +217,7 @@ class FeatureGate(L0Gate):
         self.n_features = n_features
         self.max_features = max_features
 
-    def select_features(
-        self, data: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    def select_features(self, data: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Select features from data.
 
@@ -377,9 +373,7 @@ class HybridGate(nn.Module):
 
         # Track selection type
         # Create initial selection type array
-        _ = np.array(
-            ["l0"] * len(l0_indices) + ["random"] * len(random_indices)
-        )
+        _ = np.array(["l0"] * len(l0_indices) + ["random"] * len(random_indices))
 
         # Reorder selection_type to match sorted indices
         _ = torch.cat(

l0/layers.py

@@ -222,9 +222,7 @@ class L0Conv2d(nn.Module):
             gates = gates.view(self.weight.shape)
             masked_weight = self.weight * gates
 
-        return F.conv2d(
-            input, masked_weight, self.bias, self.stride, self.padding
-        )
+        return F.conv2d(input, masked_weight, self.bias, self.stride, self.padding)
 
     def get_l0_penalty(self) -> torch.Tensor:
         """Get L0 penalty."""
@@ -398,9 +396,7 @@ class SparseMLP(nn.Module):
         """Get sparsity statistics for all layers."""
         stats = {}
         for name, module in self.named_modules():
-            if hasattr(module, "get_sparsity") and hasattr(
-                module, "get_l0_penalty"
-            ):
+            if hasattr(module, "get_sparsity") and hasattr(module, "get_l0_penalty"):
                 stats[name] = {
                     "sparsity": module.get_sparsity(),
                     "active_params": module.get_l0_penalty().item(),

l0/penalties.py

@@ -107,9 +107,7 @@ def get_sparsity_stats(model: nn.Module) -> dict[str, dict[str, Any]]:
     stats = {}
 
     for name, module in model.named_modules():
-        if hasattr(module, "get_sparsity") and hasattr(
-            module, "get_l0_penalty"
-        ):
+        if hasattr(module, "get_sparsity") and hasattr(module, "get_l0_penalty"):
             active_params = module.get_l0_penalty().item()
 
             # Calculate total parameters
@@ -204,16 +202,12 @@ class TemperatureScheduler:
         progress = epoch / self.anneal_epochs
 
         if self.schedule == "linear":
-            temp = (
-                self.initial_temp
-                - (self.initial_temp - self.final_temp) * progress
-            )
+            temp = self.initial_temp - (self.initial_temp - self.final_temp) * progress
 
         elif self.schedule == "exponential":
             log_temp = (
                 np.log(self.initial_temp)
-                - (np.log(self.initial_temp) - np.log(self.final_temp))
-                * progress
+                - (np.log(self.initial_temp) - np.log(self.final_temp)) * progress
             )
             temp = np.exp(log_temp)
 

l0/sparse.py

@@ -7,7 +7,6 @@ sparse input matrices without converting them to dense format.
 
 # ruff: noqa: N803 N806
 
-
 import numpy as np
 import torch
 import torch.nn as nn
@@ -37,6 +36,12 @@ class SparseL0Linear(nn.Module):
         Initial probability of keeping each feature
     device : str or torch.device
         Device to run computations on ('cpu' or 'cuda')
+    seed : int, optional
+        Seed for the RNG used by the ``log_alpha`` init. When set, two
+        models with the same inputs produce byte-identical initial
+        ``log_alpha`` without the caller having to manage PyTorch's
+        global RNG. ``None`` (default) preserves legacy behaviour of
+        using the global RNG.
     """
 
     def __init__(
@@ -48,6 +53,7 @@ class SparseL0Linear(nn.Module):
         zeta: float = 1.1,
         init_keep_prob: float = 0.5,
         device: str | torch.device = "cpu",
+        seed: int | None = None,
     ):
         super().__init__()
         self.n_features = n_features
@@ -56,6 +62,15 @@ class SparseL0Linear(nn.Module):
         self.gamma = gamma
         self.zeta = zeta
         self.device = torch.device(device)
+        self.seed = seed
+
+        if seed is not None:
+            self._generator: torch.Generator | None = torch.Generator(
+                device=self.device
+            )
+            self._generator.manual_seed(int(seed))
+        else:
+            self._generator = None
 
         # Model parameters
         self.weight = nn.Parameter(torch.zeros(n_features, device=self.device))
@@ -64,21 +79,16 @@ class SparseL0Linear(nn.Module):
         else:
             self.register_parameter("bias", None)
 
-        # L0 gate parameters
+        # L0 gate parameters: mu + N(0, 0.01) jitter, optionally seeded.
         mu = torch.log(torch.tensor(init_keep_prob / (1 - init_keep_prob)))
-        self.log_alpha = nn.Parameter(
-            torch.normal(
-                mu.item(), 0.01, size=(n_features,), device=self.device
-            )
-        )
+        noise = torch.randn(n_features, generator=self._generator, device=self.device)
+        self.log_alpha = nn.Parameter(mu.item() + 0.01 * noise)
 
         # Cache for sparse tensor conversion
         self._cached_X_torch: torch.sparse.Tensor | None = None
         self._cached_X_shape: tuple[int, int] | None = None
 
-    def _convert_sparse_to_torch(
-        self, X_sparse: sp.spmatrix
-    ) -> torch.sparse.Tensor:
+    def _convert_sparse_to_torch(self, X_sparse: sp.spmatrix) -> torch.sparse.Tensor:
         """
         Convert scipy sparse matrix to torch sparse tensor.
 
@@ -92,9 +102,7 @@ class SparseL0Linear(nn.Module):
             return self._cached_X_torch
 
         X_coo = X_sparse.tocoo()
-        indices = torch.LongTensor(np.vstack([X_coo.row, X_coo.col])).to(
-            self.device
-        )
+        indices = torch.LongTensor(np.vstack([X_coo.row, X_coo.col])).to(self.device)
         values = torch.FloatTensor(X_coo.data).to(self.device)
         X_torch = torch.sparse_coo_tensor(
             indices,
@@ -110,18 +118,26 @@ class SparseL0Linear(nn.Module):
 
         return X_torch
 
+    # Bounds for `log_alpha` before scaling by `beta`. Values outside
+    # ``[-LOG_ALPHA_BOUND, LOG_ALPHA_BOUND]`` saturate the sigmoid and
+    # vanish gradients; clamping keeps sampling and deterministic gates
+    # well-defined under fp16/bf16.
+    LOG_ALPHA_BOUND: float = 20.0
+
     def _sample_gates(self) -> torch.Tensor:
         """Sample gates using Hard Concrete distribution."""
         eps = 1e-6
-        u = torch.rand_like(self.log_alpha).clamp(eps, 1 - eps)
-        X = (torch.log(u) - torch.log(1 - u) + self.log_alpha) / self.beta
+        log_alpha = self.log_alpha.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        u = torch.rand_like(log_alpha).clamp(eps, 1 - eps)
+        X = (torch.log(u) - torch.log(1 - u) + log_alpha) / self.beta
         s = torch.sigmoid(X)
         s_bar = s * (self.zeta - self.gamma) + self.gamma
         return s_bar.clamp(0, 1)
 
     def get_deterministic_gates(self) -> torch.Tensor:
         """Get deterministic gate values (for inference)."""
-        X = self.log_alpha / self.beta
+        log_alpha = self.log_alpha.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        X = log_alpha / self.beta
         s = torch.sigmoid(X)
         s_bar = s * (self.zeta - self.gamma) + self.gamma
         return s_bar.clamp(0, 1)
@@ -178,7 +194,8 @@ class SparseL0Linear(nn.Module):
         c = -self.beta * torch.log(
             torch.tensor(-self.gamma / self.zeta, device=self.device)
         )
-        pi = torch.sigmoid(self.log_alpha + c)
+        log_alpha = self.log_alpha.clamp(-self.LOG_ALPHA_BOUND, self.LOG_ALPHA_BOUND)
+        pi = torch.sigmoid(log_alpha + c)
         return pi.sum()
 
     def get_sparsity(self) -> float:
@@ -291,7 +308,7 @@ class SparseL0Linear(nn.Module):
                 with torch.no_grad():
                     n_active = (self.get_deterministic_gates() > 0.01).sum()
                     print(
-                        f"Epoch {epoch+1:4d}: "
+                        f"Epoch {epoch + 1:4d}: "
                         f"loss={loss.item():.4f}, "
                         f"mse={mse_loss.item():.4f}, "
                         f"l0={l0_loss.item():.2f}, "
@@ -300,9 +317,7 @@ class SparseL0Linear(nn.Module):
 
         return self
 
-    def predict(
-        self, X_sparse: sp.spmatrix | torch.sparse.Tensor
-    ) -> torch.Tensor:
+    def predict(self, X_sparse: sp.spmatrix | torch.sparse.Tensor) -> torch.Tensor:
         """
         Make predictions on new data.
 

{l0_python-0.6.1.dist-info → l0_python-0.6.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: l0-python
-Version: 0.6.1
+Version: 0.6.2
 Summary: L0 regularization for sparse neural networks and intelligent sampling
 Author-email: PolicyEngine <hello@policyengine.org>
 License: MIT License
@@ -45,7 +45,7 @@ Requires-Dist: scipy>=1.7.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pytest-cov>=3.0; extra == "dev"
-Requires-Dist: black>=22.0; extra == "dev"
+Requires-Dist: ruff>=0.9.0; extra == "dev"
 Requires-Dist: build>=1.3; extra == "dev"
 Requires-Dist: scikit-learn>=1.0; extra == "dev"
 Requires-Dist: yaml-changelog>=0.3.0; extra == "dev"

l0_python-0.6.2.dist-info/RECORD

@@ -0,0 +1,12 @@
+l0/__init__.py,sha256=ByyYJoiHgx0c32P0VxSaV8gMWnEaWsNlr69Nu9ydgi0,1292
+l0/calibration.py,sha256=U_w_oteOKMEavagMy2TiR8binHZbRhgbOvt01r0LJNA,22957
+l0/distributions.py,sha256=My-d-IXwj0i-H2LZme1EUVqzV54VC7rtQLwQ5gFNx7I,7758
+l0/gates.py,sha256=D5x6M4-mVykKnaLYZSf4mG3cWCGdXacN4ry8KIepC-A,11393
+l0/layers.py,sha256=wBYcykrzVD2ZrkOpvDwlTwpglLKJwONc5M8JAC7uiNs,13733
+l0/penalties.py,sha256=h1cKZtY4-h_IpFnPzbDI-3nhwrlqmJArTGGbqEAc6is,9483
+l0/sparse.py,sha256=zKNB_2q9dyN9z7RkLTmZ29xNNyBPXMkCR1fGHKCp58E,10871
+l0_python-0.6.2.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
+l0_python-0.6.2.dist-info/METADATA,sha256=JXbdK2cxPFow8ybTrgJ3UYvhY0SpRYFqRErQgfz1BbU,7144
+l0_python-0.6.2.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+l0_python-0.6.2.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
+l0_python-0.6.2.dist-info/RECORD,,

{l0_python-0.6.1.dist-info → l0_python-0.6.2.dist-info}/WHEEL

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

l0_python-0.6.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-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.1.dist-info/licenses/LICENSE,sha256=LoyFJ8qmpwWPp5RmxFGHhdBxsemvp_BqGyiCgNAZHlk,1069
-l0_python-0.6.1.dist-info/METADATA,sha256=kf7gdRFBO8E9hIbx3w80qDG8xBP8fkOipr8wJ_Nn310,7144
-l0_python-0.6.1.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
-l0_python-0.6.1.dist-info/top_level.txt,sha256=brP2SoElgmGhNJk-qRFEQZ5UE-d1e87Tr60hSSh8n34,3
-l0_python-0.6.1.dist-info/RECORD,,