humancompatible-train 0.3.0__tar.gz → 0.3.0.1__tar.gz

{humancompatible_train-0.3.0/src/humancompatible_train.egg-info → humancompatible_train-0.3.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: humancompatible-train
-Version: 0.3.0
+Version: 0.3.0.1
 Summary: PyTorch-based package for constrained training of neural networks
 Author: Gilles Bareilles, Adam Bosak, Jana Lepsova, Jakub Marecek
 Author-email: Andrii Kliachkin <kliacand@fel.cvut.cz>
@@ -156,8 +156,6 @@ The constraint options are:
 - `equalized_odds_vec`: constraint on the Positive Rate of each group as defined by `fairret.NormLoss`;
 - `weight_norm`: constraint on the Frobenius norm of the weights and biases of each layer of the neural network.
 
-The above command will run the benchmark and save the logs and models in the `benchmark/results/[dataset]_[task]` directory; it will also plot the trajectories into a .png file saved in the same location.
-
 The benchmarking code (all of which is contained in the `benchmark` directory) is easy to parse and extend with other datasets and constraints.
 
 

{humancompatible_train-0.3.0 → humancompatible_train-0.3.0.1}/README.md

@@ -124,8 +124,6 @@ The constraint options are:
 - `equalized_odds_vec`: constraint on the Positive Rate of each group as defined by `fairret.NormLoss`;
 - `weight_norm`: constraint on the Frobenius norm of the weights and biases of each layer of the neural network.
 
-The above command will run the benchmark and save the logs and models in the `benchmark/results/[dataset]_[task]` directory; it will also plot the trajectories into a .png file saved in the same location.
-
 The benchmarking code (all of which is contained in the `benchmark` directory) is easy to parse and extend with other datasets and constraints.
 
 

{humancompatible_train-0.3.0 → humancompatible_train-0.3.0.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "humancompatible-train"
-version = "0.3.0"
+version = "0.3.0.1"
 dependencies = [
     "torch",
     "numpy",

{humancompatible_train-0.3.0 → humancompatible_train-0.3.0.1}/src/humancompatible/train/dual_optim/pbm.py

@@ -11,34 +11,35 @@ from .barrier import (
     quad_recipr_der
 )
 
-
-
-
 class PBM(Optimizer):
     def __init__(
         self,
         m: int = None,
-        mu: float = 0.3,
-        lr: float = None,
-        penalty_update: str = 'dimin',
-        pbf: str = 'quadratic_logarithmic',
-        init_duals: float | Tensor = 0.01,
-        init_penalties: float | Tensor = 10,
+        penalty_mult: float = 0.1,
+        gamma: float = 0.9,
+        delta: float = 0.9,
+        penalty_update: str = 'dimin_adapt',
         *,
-        dual_range: Tuple[float, float] = (1e-4, 100.),
-        momentum: float = 0.,
-        dampening: float = 0.,
-        penalty_range: Tuple[float, float] = (1e-1, 100.),
-        device = None
+        pbf: str = 'quadratic_logarithmic',
+        init_duals: float | Tensor = None,
+        init_penalties: float | Tensor = None,
+        dual_range: Tuple[float, float] = (0.0001, 100.),
+        penalty_range: Tuple[float, float] = (0.1, 100.),
+        device = None,
+        primal_update_process_length=1, # length of the primal update process - if =1, is the original algorithm
     ) -> None:
         """
         A wrapper over a PyTorch`Optimizer` that works on the dual maximization tasks according to the Penalty-Barrier Method rule. Creates and updates dual variables.
 
         :param m: Number of constraints (determines the number of dual variables to create)
         :type m: int
-        :param mu: Multiplier for penalty update. 
-        :type mu: float
-        :param penalty_update: Penalty update strategy; must be one of `dimin`,`dimin_dual`,`const`. Defaults to`dimin`.
+        :param penalty_mult: Multiplier for penalty update (K1 or K2). For K2 (adaptive penalty update), values close to 1 correspond to a high "momentum".
+        :type penalty_mult: float
+        :param gamma: Multiplier for dual parameter update. Values close to 1 correspond to a high "momentum".
+        :type gamma: float
+        :param delta: Violation/satisfaction parameter for penalty update; values > 1 make the penalties decrease faster on violated constraints and vice versa.
+        :type delta: float 
+        :param penalty_update: Penalty update strategy; must be one of `dimin`,`dimin_dual`,`dimin_adapt`,`const`. Defaults to`dimin_adapt`.
         :type penalty_update: str
         :param pbf: Penalty-Barrier Function to use. Must be one of `quadratic_logarithmic`,`quadratic_reciprocal`
         :type pbf: str
@@ -48,40 +49,28 @@ class PBM(Optimizer):
         :type init_penalties: float | Tensor
         :param dual_range: Safeguarding range for dual variables; they will be`clamp`-ed to this range.
         :type dual_range: Tuple[float, float]
-        :param momentum: Momentum/Smoothing factor for dual variables. Set to `0` to disable.
-        :type momentum: float
-        :param dampening: Dampening for momentum. Defaults to 1 - momentum if momentum is not 0 and 0 otherwise. Set to `0` to disable.
-        :type dampening: float
         """
 
-        # ## checks ##
-
         self.dual_range = dual_range
         self.penalty_range = penalty_range
 
-        if lr is None:
-            lr = mu
-
-        if momentum > 0 and dampening == 0:
-            dampening = momentum
-
-        params, defaults = self._init_constraint_group(m, mu, lr, penalty_update, pbf, init_duals, init_penalties, momentum, dampening, dual_range, penalty_range, device)
-
+        params, defaults = self._init_constraint_group(m, penalty_mult, penalty_update, delta, pbf, init_duals, init_penalties, gamma, dual_range, penalty_range, primal_update_process_length, device)
+        self.iter = 0
         super().__init__(params, defaults)
 
     @staticmethod
     def _init_constraint_group(
         m: int,
-        mu: float = None,
-        lr: float = None,
+        p_mult: float = None,
         penalty_update: str = None,
+        delta: float = None,
         pbf: str = None,
         init_duals: float | Tensor = None,
         init_penalties: float | Tensor = None,
-        momentum: float = None,
-        dampening: float = None,
+        dual_momentum: float = None,
         dual_range: Tuple[float, float] = None,
         penalty_range: Tuple[float, float] = None,
+        primal_update_process_length: int = 1,
         device = None
     ):
         if init_duals is None and m is None:
@@ -95,23 +84,29 @@ class PBM(Optimizer):
         duals = Parameter(init_duals, requires_grad=False)
         penalties = Parameter(init_penalties, requires_grad=False)
         
+        primal_update_process_length = primal_update_process_length
+
         if penalty_update == 'dimin':
             penalty_update_f = _update_penalties_dimin
         elif penalty_update == 'dimin_dual':
             penalty_update_f = _update_penalties_dimin_dual
+        elif penalty_update == 'dimin_adapt':
+            penalty_update_f = _update_penalties_adapt
         elif penalty_update == 'const':
             penalty_update_f = _update_penalties_const
-        else:
+        elif penalty_update is None:
             penalty_update_f = None
+        else:
+            raise ValueError(f'Unknown penalty update function: {penalty_update}!')
 
         settings_dict = {
-            "mu": mu,
-            "lr": lr,
+            "p_mult": p_mult,
             "penalty_update": penalty_update_f,
+            "delta": delta,
             "pbf": pbf,
-            "momentum": momentum,
-            "dampening": dampening,
-            "momentum_buffer": torch.zeros_like(init_duals, requires_grad = False, device=device),
+            "dual_momentum": dual_momentum,
+            "primal_update_process_length": primal_update_process_length,
+            "dual_momentum_buffer": torch.zeros_like(init_duals, requires_grad = False, device=device),
         }
         settings_dict = {k:v for k,v in settings_dict.items() if v is not None}
 
@@ -122,6 +117,8 @@ class PBM(Optimizer):
     @property
     def duals(self) -> Tensor:
         """
+        Returns all dual variables concatenated from all constraint groups.
+
         :return: Dual variables, concatenated into a single tensor.
         :rtype: Tensor
         """
@@ -130,6 +127,8 @@ class PBM(Optimizer):
     @property
     def penalties(self) -> Tensor:
         """
+        Returns all penalty variables concatenated from all constraint groups.
+
         :return: Penalties, concatenated into a single tensor.
         :rtype: Tensor
         """
@@ -138,35 +137,41 @@ class PBM(Optimizer):
     def add_constraint_group(
         self,
         m: int,
-        mu: float = None,
-        lr: float = None,
+        penalty_mult: float = None,
         penalty_update: str = None,
+        delta: float = None,
         pbf: str = None,
         init_duals: float | Tensor = None,
         init_penalties: float | Tensor = None,
         *,
         momentum: float = None,
-        dampening: float = None,
+        primal_update_process_length: int = 1
     ) -> None:
         """
-        Allows to add a group of dual variables with separate initial values and learning rates.
+        Adds an additional group of dual variables with separate hyperparameters and barrier functions.
 
-        :param m: Size of group (number of dual variables to add)
+        :param m: Number of constraints in this group (determines the number of dual variables to add)
         :type m: int
-        :param mu: Multiplier for penalty update. 
-        :type mu: float
-        :param penalty_update: Penalty update strategy; must be one of `dimin`,`dimin_dual`,`const`. Defaults to`dimin`.
+        :param penalty_mult: Multiplier for penalty update (K1 or K2). If None, inherits from parent. For adaptive penalty update, values close to 1 correspond to high "momentum".
+        :type penalty_mult: float
+        :param penalty_update: Penalty update strategy; must be one of `dimin`, `dimin_dual`, `dimin_adapt`, `const`. If None, defaults to `dimin`.
         :type penalty_update: str
-        :param pbf: Penalty-Barrier Function to use. Can be one of `"quadratic_logarithmic", "quadratic_reciprocal"`
+        :param delta: Violation/satisfaction parameter for penalty update. If None, inherits from parent.
+        :type delta: float
+        :param pbf: Penalty-Barrier Function to use. Must be one of `quadratic_logarithmic`, `quadratic_reciprocal`.
         :type pbf: str
-        :param init_duals: Initial values for the dual variables. Defaults to the lower bound for all.
+        :param init_duals: Initial values for the dual variables in this group. Defaults to dual lower bound for all.
         :type init_duals: float | Tensor
-        :param init_penalties: Initial values for the penalty variables. Defaults to min(10, lower_bound) for all.
+        :param init_penalties: Initial values for the penalty variables in this group. Defaults to penalty upper bound for all.
         :type init_penalties: float | Tensor
+        :param momentum: Multiplier for dual parameter update in this group. Values close to 1 correspond to high "momentum". If None, inherits from parent.
+        :type momentum: float
+        :param primal_update_process_length: Length of the primal update process for this group. If 1 (default), uses original algorithm variant.
+        :type primal_update_process_length: int
         """
         
         
-        params, settings_dict = self._init_constraint_group(m, mu, lr, penalty_update, pbf, init_duals, init_penalties, momentum, dampening, self.dual_range, self.penalty_range)
+        params, settings_dict = self._init_constraint_group(m, penalty_mult, penalty_update, delta, pbf, init_duals, init_penalties, momentum, self.dual_range, self.penalty_range, primal_update_process_length)
         param_group_dict = {"params": params, **settings_dict}
         self.add_param_group(param_group_dict)
 
@@ -181,14 +186,14 @@ class PBM(Optimizer):
         """
 
         for i, group in enumerate(self.param_groups):
-            duals, penalties, mu, lr, penalty_update, pbf, momentum, dampening, buffer = group["params"][0], group["params"][1], group["mu"], group["lr"], group["penalty_update"], group["pbf"], group['momentum'], group['dampening'], group['momentum_buffer']
+            duals, penalties, p_mult, _update_penalties, delta, pbf, momentum, primal_update_process_length  = group["params"][0], group["params"][1], group["p_mult"], group["penalty_update"], group["delta"], group["pbf"], group['dual_momentum'], group["primal_update_process_length"]
             group_constraints = constraints[i * len(duals) : (i + 1) * len(duals)]
             cdivp = group_constraints.div(penalties)
             with torch.no_grad():
-                _update_duals(duals, cdivp, penalty_barrier_funcs[pbf]['d'], mu, momentum, dampening, buffer)
+                _update_duals(duals, cdivp, penalty_barrier_funcs[pbf]['d'], momentum)
                 clamp_(duals, min=self.dual_range[0], max=self.dual_range[1])
-                # penalty_update(penalties, lr, duals)
-                # clamp_(penalties, min=self.penalty_range[0], max=self.penalty_range[1])
+                _update_penalties(penalties, p_mult, duals, penalty_barrier_funcs[pbf]['d'](group_constraints), delta)
+                clamp_(penalties, min=self.penalty_range[0], max=self.penalty_range[1])
 
     def forward(self, loss: Tensor, constraints: Tensor) -> Tensor:
         """
@@ -221,7 +226,7 @@ class PBM(Optimizer):
         Evaluates the Penalty-Barrier Lagrangian and updates the dual variables and penalties.
 
         Combines the computation of the Lagrangian and the update of dual variables and penalties
-        in a single step, which can be more efficient than calling `forward` and `update` separately.
+        in a single step.
 
         :param loss: Loss (objective function) value.
         :type loss: torch.Tensor
@@ -233,45 +238,71 @@ class PBM(Optimizer):
 
         lagrangian = torch.zeros_like(loss)
         lagrangian.add_(loss)
+        _last_c_group_index = 0
         for i, group in enumerate(self.param_groups):
-            duals, penalties, mu, lr, _update_penalties, pbf, momentum, dampening, buffer = group["params"][0], group["params"][1], group["mu"], group["lr"], group["penalty_update"], group["pbf"], group['momentum'], group['dampening'], group['momentum_buffer']
-            group_constraints = constraints[i * len(duals) : (i + 1) * len(duals)]
+            duals, penalties, p_mult, _update_penalties, delta, pbf, momentum, primal_update_process_length  = group["params"][0], group["params"][1], group["p_mult"], group["penalty_update"], group["delta"], group["pbf"], group['dual_momentum'], group["primal_update_process_length"]
+            group_constraints = constraints[_last_c_group_index : _last_c_group_index + len(duals)]
+            _last_c_group_index = _last_c_group_index + len(duals)
             # calculate lagrangian
-            cdivp = group_constraints.div(penalties)
-            # update duals and penalties
-            with torch.no_grad():
-                _update_duals(duals, cdivp, penalty_barrier_funcs[pbf]['d'], mu, momentum, dampening, buffer)
-                clamp_(duals, min=self.dual_range[0], max=self.dual_range[1])
-                # _update_penalties(penalties, lr, duals)
-                # clamp_(penalties, min=self.penalty_range[0], max=self.penalty_range[1])
+            if self.iter + 1 == primal_update_process_length: # this enables a second variant of the algorithm
+                # update duals and penalties
+                cdivp = group_constraints.div(penalties)
+                with torch.no_grad():
+                    _update_duals(duals, cdivp, penalty_barrier_funcs[pbf]['d'], momentum)
+                    clamp_(duals, min=self.dual_range[0], max=self.dual_range[1])
+                    _update_penalties(penalties, p_mult, duals, penalty_barrier_funcs[pbf]['d'](group_constraints), delta)
+                    clamp_(penalties, min=self.penalty_range[0], max=self.penalty_range[1])
 
             cdivp = group_constraints.div(penalties)
             pbf_val = penalty_barrier_funcs[pbf]['f'](cdivp)
             lagrangian.add_(duals.mul(penalties) @ pbf_val)
 
+        # update the iter
+        self.iter = (self.iter + 1) % primal_update_process_length
+
         return lagrangian
     
-
-    def update_penalties(self):
+    def update_penalties(self, constraints: Tensor) -> None:
         """
-        Performs the penalty update according to`penalty_update`.
+        Updates penalties according to the specified penalty update strategy for each constraint group.
+
+        :param constraints: Tensor of constraint violations.
+        :type constraints: torch.Tensor
+        :return: None
+        :rtype: None
         """
-        for group in self.param_groups:
-            duals, penalties, lr, _update_penalties = group["params"][0], group["params"][1], group["lr"], group["penalty_update"]
-            _update_penalties(penalties, lr, duals)
+        for i, group in enumerate(self.param_groups):
+            duals, penalties, p_mult, _update_penalties, pbf = group["params"][0], group["params"][1], group["p_mult"], group["penalty_update"], group['pbf']
+            group_constraints = constraints[i * len(duals) : (i + 1) * len(duals)]
+            _update_penalties(penalties, p_mult, duals, penalty_barrier_funcs[pbf]['d'](group_constraints))
             clamp_(penalties, min=self.penalty_range[0], max=self.penalty_range[1])
 
+
     def state_dict(self) -> dict[str, Any]:
+        """
+        Returns the state of the optimizer as a dictionary, including dual and penalty ranges and all constraint groups.
+
+        :return: Dictionary containing optimizer state with param groups and configuration.
+        :rtype: dict[str, Any]
+        """
         
         state_dict = super().state_dict()
         state_dict["state"]["penalty_range"] = self.penalty_range
         state_dict["state"]["dual_range"] = self.dual_range
         # save params themselves in state_dict instead of param ID in default PyTorch
         for id_pg, pg in enumerate(state_dict['param_groups']):
-            pg['params'] = [self.param_groups[id_pg]['params'][param_id] for param_id in pg['params'] ]
+            pg['params'] = self.param_groups[id_pg]['params']
         return state_dict
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """
+        Loads the optimizer state from a dictionary, including ranges and all constraint groups.
+
+        :param state_dict: Dictionary containing optimizer state (as returned by state_dict).
+        :type state_dict: dict[str, Any]
+        :return: None
+        :rtype: None
+        """
         self.penalty_range = state_dict["state"]["penalty_range"]
         self.dual_range = state_dict["state"]["dual_range"]
         params = state_dict["param_groups"]
@@ -286,22 +317,29 @@ penalty_barrier_funcs = {
     'quadratic_reciprocal': {'f': quad_recipr, 'd': quad_recipr_der}
 }
 
-def _update_duals(duals: Tensor, cdivp: Tensor, pbf_der: Callable, mu: float, momentum: float, dampening: float, buffer: Tensor) -> None:
+def _update_duals(duals: Tensor, cdivp: Tensor, pbf_der: Callable, gamma: float) -> None:
     pbf_der_val = pbf_der(cdivp)
-    if momentum == 0 or not buffer.any():
-        buffer = pbf_der_val
-    else:
-        buffer.mul_(momentum).add_(pbf_der_val, alpha = 1 - dampening)
+    upd = pbf_der_val.mul(duals)
+    duals.mul_(gamma).add_(upd, alpha=1-gamma)
 
-    buffer.clamp_(mu, 1/mu)
-    duals.mul_(buffer)
-    
-
-def _update_penalties_const(penalties: Tensor, mu: Tensor, duals: Tensor):
+def _update_penalties_const(penalties: Tensor, p_mult: Tensor = None, duals: Tensor = None, phi_der: Tensor = None, delta: float = None):
     pass
 
-def _update_penalties_dimin(penalties: Tensor, mu: Tensor, duals: Tensor):
-    penalties.mul_(mu)
+def _update_penalties_dimin(penalties: Tensor, p_mult: Tensor, duals: Tensor = None, phi_der: Tensor = None, delta: float = None):
+    penalties.mul_(p_mult)
+
+def _update_penalties_adapt(penalties: Tensor, p_mult: Tensor, duals: Tensor, phi_der: Tensor, delta: float):
+    d_phd = torch.where(phi_der < 1., phi_der, torch.clamp(delta * phi_der, min=1.0))
+    b = (1-p_mult)*penalties/(d_phd + 1e-8)
+    penalties.mul_(p_mult).add_(b)
+
+def _update_penalties_dimin_dual(penalties: Tensor, p_mult: Tensor, duals: Tensor, phi_der: Tensor = None, delta: float = None):
+    penalties.mul_(p_mult).mul_(duals)
+
 
-def _update_penalties_dimin_dual(penalties: Tensor, mu: Tensor, duals: Tensor):
-    penalties.mul_(mu).mul_(duals)
+penalty_update_funcs = {
+    'const': _update_penalties_const,
+    'dimin': _update_penalties_dimin,
+    'adapt': _update_penalties_adapt,
+    'dimin_dual': _update_penalties_dimin_dual
+}

{humancompatible_train-0.3.0 → humancompatible_train-0.3.0.1}/src/humancompatible/train/optim/PBM.py

@@ -556,7 +556,7 @@ class PBM(Optimizer):
             # define the augmented F and backpropagate
             F_loss = loss + self._dual_vars @ self.constraints
 
-        # TODO: implement smoothing here
+        # TODO: implement smoothing herew
 
         # clean the gradients
         # self.constraints = torch.zeros(self.m, device=self.device)

{humancompatible_train-0.3.0 → humancompatible_train-0.3.0.1/src/humancompatible_train.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: humancompatible-train
-Version: 0.3.0
+Version: 0.3.0.1
 Summary: PyTorch-based package for constrained training of neural networks
 Author: Gilles Bareilles, Adam Bosak, Jana Lepsova, Jakub Marecek
 Author-email: Andrii Kliachkin <kliacand@fel.cvut.cz>
@@ -156,8 +156,6 @@ The constraint options are:
 - `equalized_odds_vec`: constraint on the Positive Rate of each group as defined by `fairret.NormLoss`;
 - `weight_norm`: constraint on the Frobenius norm of the weights and biases of each layer of the neural network.
 
-The above command will run the benchmark and save the logs and models in the `benchmark/results/[dataset]_[task]` directory; it will also plot the trajectories into a .png file saved in the same location.
-
 The benchmarking code (all of which is contained in the `benchmark` directory) is easy to parse and extend with other datasets and constraints.
 
 

humancompatible_train-0.3.0.1/tests/test_pbm.py

@@ -0,0 +1,532 @@
+import unittest
+import torch
+from humancompatible.train.dual_optim import PBM
+from humancompatible.train.dual_optim.barrier import quad_log, quad_log_der, quad_recipr, quad_recipr_der
+
+
+class TestPBMInitialization(unittest.TestCase):
+    """Test PBM initialization and error handling."""
+
+    def test_init_with_m(self):
+        """Test initialization with m parameter."""
+        pbm = PBM(m=5, penalty_mult=0.1)
+        self.assertEqual(len(pbm.duals), 5)
+        self.assertEqual(len(pbm.penalties), 5)
+
+    def test_init_with_scalar_duals(self):
+        """Test initialization with scalar init_duals."""
+        pbm = PBM(m=3, init_duals=0.5, init_penalties=10.0)
+        self.assertTrue(torch.allclose(pbm.duals, torch.tensor([0.5, 0.5, 0.5])))
+        self.assertTrue(torch.allclose(pbm.penalties, torch.tensor([10.0, 10.0, 10.0])))
+
+    def test_init_with_tensor_duals(self):
+        """Test initialization with tensor init_duals."""
+        init_duals = torch.tensor([0.1, 0.2, 0.3])
+        pbm = PBM(m=3, init_duals=init_duals, init_penalties=2.0)
+        self.assertTrue(torch.allclose(pbm.duals, init_duals))
+
+    def test_init_defaults_to_dual_range(self):
+        """Test that None init_duals defaults to dual_range lower bound."""
+        pbm = PBM(m=3, dual_range=(0.01, 100.0))
+        self.assertTrue(torch.allclose(pbm.duals, torch.tensor([0.01, 0.01, 0.01])))
+
+    def test_init_defaults_to_penalty_range(self):
+        """Test that None init_penalties defaults to penalty_range upper bound."""
+        pbm = PBM(m=3, penalty_range=(0.1, 50.0))
+        self.assertTrue(torch.allclose(pbm.penalties, torch.tensor([50.0, 50.0, 50.0])))
+
+    def test_init_missing_m_and_duals_raises_error(self):
+        """Test that ValueError is raised when both m and init_duals are None."""
+        with self.assertRaises(ValueError):
+            PBM(m=None, init_duals=None)
+
+    def test_init_with_device(self):
+        """Test initialization on specific device."""
+        device = 'cpu'
+        pbm = PBM(m=3, device=device)
+        self.assertEqual(pbm.duals.device.type, device)
+
+    def test_init_different_pbf(self):
+        """Test initialization with different penalty-barrier functions."""
+        pbm_log = PBM(m=3, pbf='quadratic_logarithmic')
+        pbm_recipr = PBM(m=3, pbf='quadratic_reciprocal')
+        # Just verify they don't raise errors
+        self.assertEqual(len(pbm_log.duals), 3)
+        self.assertEqual(len(pbm_recipr.duals), 3)
+
+    def test_init_invalid_pbf_raises_error(self):
+        """Test that invalid pbf raises KeyError when used."""
+        pbm = PBM(m=3, pbf='invalid_function')
+        # PBM creation doesn't validate pbf, but using it will fail
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.1, 0.1, 0.1])
+        with self.assertRaises(KeyError):
+            pbm.forward(loss, constraints)
+
+    def test_init_invalid_penalty_update_raises_error(self):
+        """Test that invalid penalty_update raises ValueError."""
+        with self.assertRaises(ValueError):
+            PBM(m=3, penalty_update='invalid_strategy')
+
+
+class TestDualAndPenaltyBounds(unittest.TestCase):
+    """Test that duals and penalties stay within specified bounds."""
+
+    def test_duals_clamped_to_range(self):
+        """Test that duals are clamped to dual_range."""
+        pbm = PBM(m=2, dual_range=(0.01, 10.0), init_duals=5.0)
+        # Large constraint violations through forward_update
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([1000.0, 1000.0])
+        for _ in range(5):
+            pbm.forward_update(loss, constraints)
+        # Duals should not exceed upper bound
+        self.assertTrue(torch.all(pbm.duals <= 10.0))
+        # Duals should not fall below lower bound
+        self.assertTrue(torch.all(pbm.duals >= 0.01))
+
+    def test_penalties_clamped_to_range(self):
+        """Test that penalties are clamped to penalty_range."""
+        pbm = PBM(m=2, penalty_mult=0.05, penalty_range=(0.1, 100.0))
+        loss = torch.tensor(1.0)
+        # Run multiple updates
+        for _ in range(10):
+            constraints = torch.tensor([0.5, 0.5])
+            pbm.forward_update(loss, constraints)
+        # Penalties should not exceed upper bound
+        self.assertTrue(torch.all(pbm.penalties <= 100.0))
+        # Penalties should not fall below lower bound
+        self.assertTrue(torch.all(pbm.penalties >= 0.1))
+
+    def test_duals_stay_positive_with_small_range(self):
+        """Test duals with very small range."""
+        pbm = PBM(m=3, dual_range=(1e-6, 1e-3), init_duals=5e-4)
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.1, 0.2, 0.3])
+        pbm.forward_update(loss, constraints)
+        self.assertTrue(torch.all(pbm.duals >= 1e-6))
+        self.assertTrue(torch.all(pbm.duals <= 1e-3))
+
+
+class TestPenaltyUpdateStrategies(unittest.TestCase):
+    """Test different penalty update strategies."""
+
+    def test_constant_penalty(self):
+        """Test that penalty_update='const' keeps penalties constant during forward_update."""
+        init_penalty = 5.0
+        pbm = PBM(m=3, penalty_update='const', init_penalties=init_penalty)
+        penalties_before = pbm.penalties.clone()
+        loss = torch.tensor(1.0)
+        
+        for _ in range(5):
+            constraints = torch.tensor([0.1, 0.2, 0.3])
+            pbm.forward_update(loss, constraints)
+        
+        self.assertTrue(torch.allclose(pbm.penalties, penalties_before))
+
+    def test_diminishing_penalty_monotonic(self):
+        """Test that penalty_update='dimin' decreases penalties monotonically."""
+        penalty_mult = 0.8
+        pbm = PBM(m=2, penalty_update='dimin', penalty_mult=penalty_mult, init_penalties=100.0)
+        
+        penalties_history = [pbm.penalties.clone()]
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.5, 0.5])
+        
+        for _ in range(5):
+            pbm.forward_update(loss, constraints)
+            penalties_history.append(pbm.penalties.clone())
+        
+        # Verify monotonic decrease (from 1 onwards since 0 is initial)
+        for i in range(len(penalties_history) - 1):
+            self.assertTrue(torch.all(penalties_history[i + 1] <= penalties_history[i]))
+
+    def test_adaptive_penalty_responds_to_violations(self):
+        """Test that adaptive penalty responds to constraint violations."""
+        pbm_adapt = PBM(m=1, penalty_update='dimin_adapt', penalty_mult=0.8, delta=0.9, init_penalties=10.0)
+        pbm_const = PBM(m=1, penalty_update='const', init_penalties=10.0)
+        loss = torch.tensor(1.0)
+        
+        # Use larger constraint to see penalty difference
+        constraint = torch.tensor([1.0])
+        pbm_adapt.forward_update(loss, constraint)
+        pbm_const.forward_update(loss, constraint)
+        
+        # Adaptive should differ from constant after update
+        # (or at least both updates should complete without error)
+        self.assertEqual(len(pbm_adapt.penalties), 1)
+
+    # def test_dimin_dual_penalty_changes(self):
+    #     """Test that dimin_dual strategy multiplies by duals."""
+    #     pbm = PBM(m=2, penalty_update='dimin_dual', penalty_mult=0.8, init_duals=2.0, init_penalties=10.0)
+        
+    #     initial_penalty = pbm.penalties.clone()
+    #     loss = torch.tensor(1.0)
+    #     # Use larger constraints to trigger meaningful updates
+    #     constraints = torch.tensor([1.0, 1.0])
+    #     pbm.forward_update(loss, constraints)
+        
+    #     # With dimin_dual, penalties should decrease due to low initial duals
+    #     breakpoint()
+    #     self.assertTrue(torch.all(pbm.penalties <= initial_penalty))
+
+
+class TestBarrierFunctions(unittest.TestCase):
+    """Test barrier function properties."""
+
+    def test_quad_log_positive_derivative_at_violation(self):
+        """Test that quad_log derivative is positive at constraint violations."""
+        t = torch.tensor([-0.2, 0.0, 0.5])
+        deriv = quad_log_der(t)
+        self.assertTrue(torch.all(deriv > 0.0))
+
+    def test_quad_recipr_positive_derivative_at_violation(self):
+        """Test that quad_recipr derivative is positive at constraint violations."""
+        t = torch.tensor([-0.2, 0.0, 0.3])
+        deriv = quad_recipr_der(t)
+        self.assertTrue(torch.all(deriv > 0.0))
+
+    def test_quad_log_derivative_matches_finite_diff_smooth_region(self):
+        """Test quad_log derivative against finite differences in smooth region."""
+        # Use values well within smooth region (t >= -0.5)
+        t = torch.tensor([0.5, 1.0, 2.0])
+        eps = 1e-3
+        
+        # Finite difference
+        deriv_fd = (quad_log(t + eps) - quad_log(t - eps)) / (2 * eps)
+        deriv_analytical = quad_log_der(t)
+        
+        self.assertTrue(torch.allclose(deriv_fd, deriv_analytical, atol=1e-3))
+
+    def test_barrier_functions_increase_with_violation(self):
+        """Test that barrier functions penalize larger violations more."""
+        t1 = torch.tensor([0.0])
+        t2 = torch.tensor([0.5])
+        
+        log_v1 = quad_log(t1)
+        log_v2 = quad_log(t2)
+        self.assertTrue(log_v2 > log_v1)
+        
+        recipr_v1 = quad_recipr(t1)
+        recipr_v2 = quad_recipr(t2)
+        self.assertTrue(recipr_v2 > recipr_v1)
+
+    def test_barrier_functions_nonzero_at_zero(self):
+        """Test that barrier function values at zero constraint are reasonable."""
+        t = torch.tensor([0.0])
+        log_val = quad_log(t)
+        recipr_val = quad_recipr(t)
+        
+        # Both should be exactly 0 at t=0
+        self.assertTrue(torch.allclose(log_val, torch.tensor([0.0]), atol=1e-6))
+        self.assertTrue(torch.allclose(recipr_val, torch.tensor([0.0]), atol=1e-6))
+
+
+class TestDualVariableUpdates(unittest.TestCase):
+    """Test dual variable update mechanisms."""
+
+    def test_dual_update_with_zero_momentum(self):
+        """Test dual update with gamma=0 (complete replacement)."""
+        pbm = PBM(m=2, gamma=0.0, init_duals=1.0, init_penalties=10.0)
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.5, 0.5])
+        
+        # With gamma=0, new duals should be completely replaced
+        pbm.forward_update(loss, constraints)
+        # Just verify update happened without error
+        self.assertEqual(len(pbm.duals), 2)
+
+    def test_dual_update_with_high_momentum(self):
+        """Test dual update with gamma close to 1 (damped update)."""
+        pbm_high = PBM(m=2, gamma=0.99, init_duals=1.0)
+        pbm_low = PBM(m=2, gamma=0.1, init_duals=1.0)
+        loss = torch.tensor(1.0)
+        
+        constraints = torch.tensor([0.5, 0.5])
+        duals_high_before = pbm_high.duals.clone()
+        duals_low_before = pbm_low.duals.clone()
+        
+        pbm_high.forward_update(loss, constraints)
+        pbm_low.forward_update(loss, constraints)
+        
+        # High momentum should change duals less
+        high_change = torch.abs(pbm_high.duals - duals_high_before).sum()
+        low_change = torch.abs(pbm_low.duals - duals_low_before).sum()
+        self.assertTrue(high_change < low_change)
+
+    def test_dual_momentum_buffer_tracked(self):
+        """Test that momentum buffer is tracked across updates."""
+        pbm = PBM(m=2, gamma=0.8)
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.5, 0.5])
+        
+        # Run multiple updates to accumulate momentum
+        for _ in range(3):
+            pbm.forward_update(loss, constraints)
+        # Just verify no errors occur
+        self.assertEqual(len(pbm.duals), 2)
+
+
+class TestMultipleConstraintGroups(unittest.TestCase):
+    """Test adding and managing multiple constraint groups."""
+
+    def test_add_constraint_group(self):
+        """Test adding a second constraint group."""
+        pbm = PBM(m=2, penalty_mult=0.1)
+        self.assertEqual(len(pbm.duals), 2)
+        
+        pbm.add_constraint_group(m=3, penalty_mult=0.15)
+        self.assertEqual(len(pbm.duals), 5)  # 2 + 3
+        self.assertEqual(len(pbm.penalties), 5)
+
+    def test_constraint_groups_independent_updates(self):
+        """Test that constraint groups update independently."""
+        pbm = PBM(m=2, penalty_mult=0.8, penalty_update='dimin')
+        pbm.add_constraint_group(m=2, penalty_mult=0.5, penalty_update='dimin', delta=0.9, pbf='quadratic_logarithmic')
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.5, 0.5, 0.1, 0.1])
+        pbm.forward_update(loss, constraints)
+        
+        # Just verify both groups were updated
+        self.assertEqual(len(pbm.penalties), 4)
+
+    def test_multiple_groups_different_pbf(self):
+        """Test adding groups with different barrier functions."""
+        pbm = PBM(m=2, pbf='quadratic_logarithmic', penalty_mult=0.1)
+        pbm.add_constraint_group(m=2, pbf='quadratic_reciprocal', penalty_mult=0.1, penalty_update='dimin', delta=0.9)
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.5, 0.5, 0.5, 0.5])
+        pbm.forward_update(loss, constraints)
+        
+        # Just verify both groups update without error
+        self.assertEqual(len(pbm.duals), 4)
+
+    def test_duals_property_concatenates_groups(self):
+        """Test that duals property correctly concatenates groups."""
+        pbm = PBM(m=3, init_duals=1.0, penalty_mult=0.1)
+        pbm.add_constraint_group(m=2, init_duals=2.0, penalty_mult=0.1, penalty_update='dimin', delta=0.9, pbf='quadratic_logarithmic')
+        
+        duals_concat = pbm.duals
+        self.assertEqual(len(duals_concat), 5)
+        self.assertTrue(torch.allclose(duals_concat[:3], torch.tensor([1.0, 1.0, 1.0])))
+        self.assertTrue(torch.allclose(duals_concat[3:], torch.tensor([2.0, 2.0])))
+
+    def test_penalties_property_concatenates_groups(self):
+        """Test that penalties property correctly concatenates groups."""
+        pbm = PBM(m=2, init_penalties=5.0, penalty_mult=0.1)
+        pbm.add_constraint_group(m=3, init_penalties=10.0, penalty_mult=0.1, penalty_update='dimin', delta=0.9, pbf='quadratic_logarithmic')
+        
+        penalties_concat = pbm.penalties
+        self.assertEqual(len(penalties_concat), 5)
+        self.assertTrue(torch.allclose(penalties_concat[:2], torch.tensor([5.0, 5.0])))
+        self.assertTrue(torch.allclose(penalties_concat[2:], torch.tensor([10.0, 10.0, 10.0])))
+
+
+class TestLagrangianComputation(unittest.TestCase):
+    """Test Lagrangian computation accuracy."""
+
+    def test_lagrangian_includes_loss(self):
+        """Test that Lagrangian includes the loss term."""
+        pbm = PBM(m=2)
+        loss = torch.tensor(5.0)
+        constraints = torch.tensor([0.1, 0.2])
+        
+        lagrangian = pbm.forward(loss, constraints)
+        # Lagrangian should be at least the loss
+        self.assertTrue(lagrangian >= loss)
+
+    def test_lagrangian_zero_constraint_equals_loss(self):
+        """Test that Lagrangian equals loss when constraints are zero."""
+        pbm = PBM(m=2, init_duals=1.0, init_penalties=1.0)
+        loss = torch.tensor(3.0)
+        constraints = torch.tensor([0.0, 0.0])
+        
+        lagrangian = pbm.forward(loss, constraints)
+        self.assertTrue(torch.allclose(lagrangian, loss))
+
+    def test_lagrangian_increases_with_violation(self):
+        """Test that Lagrangian increases with constraint violation."""
+        pbm = PBM(m=1, init_duals=1.0, init_penalties=1.0)
+        loss = torch.tensor(1.0)
+        
+        lag_small = pbm.forward(loss, torch.tensor([0.1]))
+        lag_large = pbm.forward(loss, torch.tensor([0.5]))
+        
+        self.assertTrue(lag_large > lag_small)
+
+    def test_forward_update_produces_same_lagrangian(self):
+        """Test that forward_update produces valid Lagrangian."""
+        pbm = PBM(m=2, gamma=0.8, init_duals=0.5, init_penalties=5.0)
+        
+        loss = torch.tensor(2.0)
+        constraints = torch.tensor([0.3, 0.4])
+        
+        # forward_update should produce a Lagrangian value
+        lag = pbm.forward_update(loss, constraints)
+        
+        # Lagrangian should be a scalar
+        self.assertTrue(lag.dim() == 0 or lag.shape == torch.Size([]))
+
+    def test_manual_lagrangian_computation(self):
+        """Test Lagrangian against manual calculation with simple values."""
+        pbm = PBM(m=1, init_duals=1.0, init_penalties=1.0, pbf='quadratic_logarithmic')
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.0])
+        
+        lagrangian = pbm.forward(loss, constraints)
+        
+        # When constraint is 0, barrier function contribution should be 0
+        # So lagrangian should equal loss
+        self.assertTrue(torch.allclose(lagrangian, loss, atol=1e-5))
+
+
+class TestSerialization(unittest.TestCase):
+    """Test state dict save and load."""
+
+    def test_state_dict_contains_ranges(self):
+        """Test that state_dict includes dual and penalty ranges."""
+        pbm = PBM(m=3, dual_range=(0.01, 50.0), penalty_range=(0.5, 200.0))
+        state = pbm.state_dict()
+        
+        self.assertEqual(state["state"]["dual_range"], (0.01, 50.0))
+        self.assertEqual(state["state"]["penalty_range"], (0.5, 200.0))
+
+    def test_state_dict_roundtrip_preserves_duals_and_penalties(self):
+        """Test that save/load preserves duals and penalties."""
+        pbm1 = PBM(m=3, init_duals=1.5, init_penalties=8.0, penalty_update='dimin')
+        loss = torch.tensor(1.0)
+        
+        # Modify state
+        constraints = torch.tensor([0.1, 0.2, 0.3])
+        pbm1.forward_update(loss, constraints)
+        
+        duals_before = pbm1.duals.clone()
+        penalties_before = pbm1.penalties.clone()
+        
+        # Save and load
+        state = pbm1.state_dict()
+        pbm2 = PBM(m=3, penalty_update='dimin')
+        pbm2.load_state_dict(state)
+        
+        self.assertTrue(torch.allclose(pbm2.duals, duals_before))
+        self.assertTrue(torch.allclose(pbm2.penalties, penalties_before))
+
+    def test_state_dict_roundtrip_with_multiple_groups(self):
+        """Test state dict roundtrip with multiple constraint groups."""
+        pbm1 = PBM(m=2, init_duals=1.0, init_penalties=5.0, penalty_mult=0.1, penalty_update='dimin')
+        pbm1.add_constraint_group(m=3, init_duals=0.5, init_penalties=10.0, penalty_mult=0.1, penalty_update='dimin', delta=0.9, pbf='quadratic_logarithmic')
+        loss = torch.tensor(1.0)
+        
+        # Update some
+        constraints = torch.tensor([0.1, 0.1, 0.2, 0.2, 0.2])
+        pbm1.forward_update(loss, constraints)
+        
+        state = pbm1.state_dict()
+        pbm2 = PBM(m=2, penalty_mult=0.1, penalty_update='dimin')
+        pbm2.add_constraint_group(m=3, penalty_mult=0.1, penalty_update='dimin', delta=0.9, pbf='quadratic_logarithmic')
+        pbm2.load_state_dict(state)
+        
+        self.assertTrue(torch.allclose(pbm2.duals, pbm1.duals))
+        self.assertTrue(torch.allclose(pbm2.penalties, pbm1.penalties))
+
+
+class TestParameterInteractions(unittest.TestCase):
+    """Test interactions between different parameters."""
+
+    def test_gamma_affects_dual_convergence_rate(self):
+        """Test that higher gamma slows dual changes."""
+        pbm_fast = PBM(m=1, gamma=0.1, init_duals=1.0, init_penalties=10.0)
+        pbm_slow = PBM(m=1, gamma=0.9, init_duals=1.0, init_penalties=10.0)
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([1.0])
+        
+        duals_fast_before = pbm_fast.duals.clone()
+        duals_slow_before = pbm_slow.duals.clone()
+        
+        pbm_fast.forward_update(loss, constraints)
+        pbm_slow.forward_update(loss, constraints)
+        
+        change_fast = torch.abs(pbm_fast.duals - duals_fast_before)
+        change_slow = torch.abs(pbm_slow.duals - duals_slow_before)
+        
+        self.assertTrue(change_fast > change_slow)
+
+    def test_delta_affects_adaptive_penalty_behavior(self):
+        """Test that delta parameter affects adaptive penalty update."""
+        pbm_low_delta = PBM(m=1, penalty_update='dimin_adapt', delta=0.5, penalty_mult=0.8, init_penalties=10.0)
+        pbm_high_delta = PBM(m=1, penalty_update='dimin_adapt', delta=2.0, penalty_mult=0.8, init_penalties=10.0)
+        loss = torch.tensor(1.0)
+        
+        # Constraint with significant violation
+        constraints = torch.tensor([2.0])
+        
+        pbm_low_delta.forward_update(loss, constraints)
+        pbm_high_delta.forward_update(loss, constraints)
+        
+        # Different delta should produce different penalties
+        self.assertFalse(torch.allclose(pbm_low_delta.penalties, pbm_high_delta.penalties))
+
+    def test_penalty_mult_scale_different_strategies(self):
+        """Test that penalty_mult works across different strategies."""
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.5])
+        
+        pbm_dimin = PBM(m=1, penalty_update='dimin', penalty_mult=0.5, init_penalties=100.0)
+        pbm_dimin.forward_update(loss, constraints)
+        dimin_penalty = pbm_dimin.penalties.clone()
+        
+        pbm_adapt = PBM(m=1, penalty_update='dimin_adapt', penalty_mult=0.5, init_penalties=100.0)
+        pbm_adapt.forward_update(loss, constraints)
+        
+        # Both should be affected by penalty_mult, though differently
+        self.assertTrue(torch.all(dimin_penalty < 100.0))
+        self.assertTrue(torch.all(pbm_adapt.penalties < 100.0))
+
+
+class TestConstraintSatisfactionProgress(unittest.TestCase):
+    """Test constraint satisfaction and convergence behavior."""
+
+    def test_constraints_drive_duals_up(self):
+        """Test that constraints drive dual variables upward."""
+        pbm = PBM(m=1, gamma=0.5, init_duals=0.1, init_penalties=10.0)
+        duals_before = pbm.duals.clone()
+        loss = torch.tensor(1.0)
+        
+        # Positive constraint violation
+        constraints = torch.tensor([1.0])
+        pbm.forward_update(loss, constraints)
+        
+        duals_after = pbm.duals.clone()
+        self.assertTrue(duals_after > duals_before)
+
+    def test_lagrangian_monotonic_on_consistent_constraints(self):
+        """Test that Lagrangian progression with consistent constraints."""
+        pbm = PBM(m=1, gamma=0.7, init_duals=0.5, init_penalties=5.0)
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.5])
+        
+        lagrangians = []
+        for _ in range(3):
+            lag = pbm.forward_update(loss, constraints)
+            lagrangians.append(lag.item())
+        
+        # Just verify sequence is computed without error
+        self.assertEqual(len(lagrangians), 3)
+
+    def test_penalties_decrease_with_dimin(self):
+        """Test that penalties monotonically decrease with dimin strategy."""
+        pbm = PBM(m=2, penalty_update='dimin', penalty_mult=0.7, init_penalties=100.0)
+        loss = torch.tensor(1.0)
+        constraints = torch.tensor([0.1, 0.1])
+        
+        penalties_list = [pbm.penalties.clone()]
+        for _ in range(5):
+            pbm.forward_update(loss, constraints)
+            penalties_list.append(pbm.penalties.clone())
+        
+        for i in range(len(penalties_list) - 1):
+            self.assertTrue(torch.all(penalties_list[i + 1] <= penalties_list[i]))
+
+
+if __name__ == '__main__':
+    unittest.main()

humancompatible_train-0.3.0/tests/test_pbm.py

@@ -1,94 +0,0 @@
-import unittest
-import torch
-from torch.optim import Optimizer
-from humancompatible.train.dual_optim.pbm import PBM
-from humancompatible.train.dual_optim.barrier import quad_log, quad_log_der, quad_recipr, quad_recipr_der
-
-class TestPBM(unittest.TestCase):
-    def setUp(self):
-        # Initialize PBM instances for reuse in tests
-        self.pbm_default = PBM(m=3, mu=0.3, lr=0.1, penalty_update='dimin', pbf='quadratic_logarithmic')
-        self.pbm_reciprocal = PBM(m=3, mu=0.3, lr=0.1, penalty_update='dimin', pbf='quadratic_reciprocal')
-        self.pbm_const_penalty = PBM(m=3, mu=0.3, lr=0.1, penalty_update='const', pbf='quadratic_logarithmic')
-        self.pbm_dimin_dual = PBM(m=3, mu=0.3, lr=0.1, penalty_update='dimin_dual', pbf='quadratic_logarithmic', init_duals=2.0)
-
-        # Common test data
-        self.loss = torch.tensor(5.0)
-        self.constraints = torch.tensor([1.0, 2.0, 3.0])
-        self.large_constraints = torch.tensor([10.0, 20.0, 30.0])
-
-    def test_pbm_initialization(self):
-        # Test initialization with m
-        self.assertEqual(len(self.pbm_default.duals), 3)
-        self.assertEqual(len(self.pbm_default.penalties), 3)
-        self.assertEqual(self.pbm_default.defaults['mu'], 0.3)
-        self.assertEqual(self.pbm_default.defaults['pbf'], 'quadratic_logarithmic')
-
-        # Test initialization with init_duals and init_penalties
-        init_duals = torch.tensor([1.0, 2.0, 3.0])
-        init_penalties = torch.tensor([10.0, 20.0, 30.0])
-        pbm = PBM(init_duals=init_duals, init_penalties=init_penalties, mu=0.3, lr=0.1, penalty_update='dimin', pbf='quadratic_logarithmic')
-        self.assertTrue(torch.all(pbm.duals == init_duals))
-        self.assertTrue(torch.all(pbm.penalties == init_penalties))
-
-        # Test invalid initialization
-        with self.assertRaises(ValueError):
-            PBM(m=None, init_duals=None)
-
-    def test_pbm_forward(self):
-        # Test Lagrangian computation
-        lagrangian = self.pbm_default.forward(self.loss, self.constraints)
-        cdivp = self.constraints / self.pbm_default.penalties
-        pbf_val = quad_log(cdivp)
-        expected_lagrangian = self.loss + torch.dot(self.pbm_default.duals * self.pbm_default.penalties, pbf_val)
-        self.assertTrue(torch.allclose(lagrangian, expected_lagrangian))
-
-    def test_pbm_update_duals(self):
-        # Test dual variable update
-        cdivp = self.constraints / self.pbm_default.penalties
-        pbf_der_val = quad_log_der(cdivp)
-        expected_duals = self.pbm_default.duals * torch.clamp(pbf_der_val, 0.3, 1/0.3)
-
-        self.pbm_default.update(self.constraints)
-        self.assertTrue(torch.allclose(self.pbm_default.duals, expected_duals))
-
-    def test_pbm_update_penalties_dimin(self):
-        # Test penalty update with 'dimin' strategy
-        initial_penalties = self.pbm_default.penalties.clone()
-        self.pbm_default.update_penalties()
-        expected_penalties = initial_penalties * 0.1
-        self.assertTrue(torch.allclose(self.pbm_default.penalties, expected_penalties))
-
-    def test_pbm_update_penalties_dimin_dual(self):
-        # Test penalty update with 'dimin_dual' strategy
-        initial_penalties = self.pbm_dimin_dual.penalties.clone()
-        expected_penalties = initial_penalties * 0.1 * self.pbm_dimin_dual.duals
-
-        self.pbm_dimin_dual.update_penalties()
-        self.assertTrue(torch.allclose(self.pbm_dimin_dual.penalties, expected_penalties))
-
-    def test_pbm_update_penalties_const(self):
-        # Test penalty update with 'const' strategy
-        initial_penalties = self.pbm_const_penalty.penalties.clone()
-        self.pbm_const_penalty.update_penalties()
-        self.assertTrue(torch.allclose(self.pbm_const_penalty.penalties, initial_penalties))
-
-    def test_pbm_forward_update(self):
-        # Test combined forward and update
-        lagrangian = self.pbm_default.forward_update(self.loss, self.constraints)
-        cdivp = self.constraints / self.pbm_default.penalties
-        pbf_val = quad_log(cdivp)
-        expected_lagrangian = self.loss + torch.dot(self.pbm_default.duals * self.pbm_default.penalties, pbf_val)
-        self.assertTrue(torch.allclose(lagrangian, expected_lagrangian))
-
-    def test_pbm_add_constraint_group(self):
-        # Test adding a new constraint group
-        self.pbm_default.add_constraint_group(m=2, mu=0.5, penalty_update='dimin', pbf='quadratic_reciprocal')
-        self.assertEqual(len(self.pbm_default.duals), 5)
-        self.assertEqual(len(self.pbm_default.penalties), 5)
-        self.assertEqual(self.pbm_default.param_groups[1]['mu'], 0.5)
-        self.assertEqual(self.pbm_default.param_groups[1]['pbf'], 'quadratic_reciprocal')
-        self.assertEqual(self.pbm_default.param_groups[1]['lr'], 0.1)
-
-if __name__ == '__main__':
-    unittest.main()