congrads 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

congrads/core/batch_runner.py

@@ -0,0 +1,200 @@
+"""Defines the BatchRunner, which executes individual batches for training, validation, and testing.
+
+Responsibilities:
+- Move batch data to the appropriate device
+- Run forward passes through the network
+- Compute base and constraint-adjusted losses
+- Perform backpropagation during training
+- Accumulate metrics for loss and other monitored quantities
+- Trigger callbacks at key points in the batch lifecycle
+"""
+
+import torch
+from torch import Tensor
+from torch.nn import Module
+from torch.optim import Optimizer
+
+from ..callbacks.base import CallbackManager
+from ..core.constraint_engine import ConstraintEngine
+from ..metrics import MetricManager
+
+
+class BatchRunner:
+    """Executes a single batch for training, validation, or testing.
+
+    The BatchRunner handles moving data to the correct device, running the network
+    forward, computing base and constraint-adjusted losses, performing backpropagation
+    during training, accumulating metrics, and dispatching callbacks at key points
+    in the batch lifecycle.
+    """
+
+    def __init__(
+        self,
+        network: Module,
+        criterion,
+        optimizer: Optimizer,
+        constraint_engine: ConstraintEngine,
+        metric_manager: MetricManager | None,
+        callback_manager: CallbackManager | None,
+        device: torch.device,
+    ):
+        """Initialize the BatchRunner.
+
+        Args:
+            network: The neural network module to execute.
+            criterion: Loss function callable accepting (output, target, data=batch).
+            optimizer: Optimizer for updating network parameters.
+            constraint_engine: ConstraintEngine instance for evaluating and enforcing constraints.
+            metric_manager: Optional MetricManager for logging batch metrics.
+            callback_manager: Optional CallbackManager for triggering hooks during batch processing.
+            device: Torch device on which to place data and network.
+        """
+        self.network = network
+        self.criterion = criterion
+        self.optimizer = optimizer
+        self.constraint_engine = constraint_engine
+        self.metric_manager = metric_manager
+        self.callback_manager = callback_manager
+        self.device = device
+
+    def _to_device(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Move all tensors in the batch to the BatchRunner's device.
+
+        Args:
+            batch: Dictionary of tensors for a single batch.
+
+        Returns:
+            Dictionary of tensors moved to the target device.
+        """
+        return {k: v.to(self.device) for k, v in batch.items()}
+
+    def _run_callbacks(self, hook: str, data: dict) -> dict:
+        """Run the specified callback hook on the batch data.
+
+        Args:
+            hook: Name of the callback hook to run.
+            data: Dictionary containing batch data.
+
+        Returns:
+            Potentially modified batch data after callback execution.
+        """
+        if self.callback_manager is None:
+            return data
+        return self.callback_manager.run(hook, data)
+
+    def train_batch(self, batch: dict[str, Tensor]) -> Tensor:
+        """Run a single training batch.
+
+        Steps performed:
+        1. Move batch to device and run "on_train_batch_start" callback.
+        2. Forward pass through the network.
+        3. Compute base loss using the criterion and accumulate metric.
+        4. Apply constraint-based adjustments to the loss.
+        5. Perform backward pass and optimizer step.
+        6. Run "on_train_batch_end" callback.
+
+        Args:
+            batch: Dictionary of input and target tensors for the batch.
+
+        Returns:
+            Tensor: The base loss computed before constraint adjustments.
+        """
+        batch = self._to_device(batch)
+        batch = self._run_callbacks("on_train_batch_start", batch)
+
+        # Forward
+        batch = self.network(batch)
+        batch = self._run_callbacks("after_train_forward", batch)
+
+        # Base loss
+        loss: Tensor = self.criterion(
+            batch["output"],
+            batch["target"],
+            data=batch,
+        )
+
+        if self.metric_manager is not None:
+            self.metric_manager.accumulate("Loss/train", loss.unsqueeze(0))
+
+        # Constraint-adjusted loss
+        combined_loss = self.constraint_engine.train(batch, loss)
+
+        # Backward
+        self.optimizer.zero_grad()
+        combined_loss.backward()
+        self.optimizer.step()
+
+        batch = self._run_callbacks("on_train_batch_end", batch)
+        return loss
+
+    def valid_batch(self, batch: dict[str, Tensor]) -> Tensor:
+        """Run a single validation batch.
+
+        Steps performed:
+        1. Move batch to device and run "on_valid_batch_start" callback.
+        2. Forward pass through the network.
+        3. Compute base loss using the criterion and accumulate metric.
+        4. Evaluate constraints via the ConstraintEngine (does not modify loss).
+        5. Run "on_valid_batch_end" callback.
+
+        Args:
+            batch: Dictionary of input and target tensors for the batch.
+
+        Returns:
+            Tensor: The base loss computed for the batch.
+        """
+        batch = self._to_device(batch)
+        batch = self._run_callbacks("on_valid_batch_start", batch)
+
+        batch = self.network(batch)
+        batch = self._run_callbacks("after_valid_forward", batch)
+
+        loss: Tensor = self.criterion(
+            batch["output"],
+            batch["target"],
+            data=batch,
+        )
+
+        if self.metric_manager is not None:
+            self.metric_manager.accumulate("Loss/valid", loss.unsqueeze(0))
+
+        self.constraint_engine.validate(batch, loss)
+
+        batch = self._run_callbacks("on_valid_batch_end", batch)
+        return loss
+
+    def test_batch(self, batch: dict[str, Tensor]) -> Tensor:
+        """Run a single test batch.
+
+        Steps performed:
+        1. Move batch to device and run "on_test_batch_start" callback.
+        2. Forward pass through the network.
+        3. Compute base loss using the criterion and accumulate metric.
+        4. Evaluate constraints via the ConstraintEngine (does not modify loss).
+        5. Run "on_test_batch_end" callback.
+
+        Args:
+            batch: Dictionary of input and target tensors for the batch.
+
+        Returns:
+            Tensor: The base loss computed for the batch.
+        """
+        batch = self._to_device(batch)
+        batch = self._run_callbacks("on_test_batch_start", batch)
+
+        batch = self.network(batch)
+        batch = self._run_callbacks("after_test_forward", batch)
+
+        loss: Tensor = self.criterion(
+            batch["output"],
+            batch["target"],
+            data=batch,
+        )
+
+        if self.metric_manager is not None:
+            self.metric_manager.accumulate("Loss/test", loss.unsqueeze(0))
+
+        self.constraint_engine.test(batch, loss)
+
+        batch = self._run_callbacks("on_test_batch_end", batch)
+        return loss

congrads/core/congradscore.py

@@ -0,0 +1,271 @@
+"""This module provides the core CongradsCore class for the main training functionality.
+
+It is designed to integrate constraint-guided optimization into neural network training.
+It extends traditional training processes by enforcing specific constraints
+on the model's outputs, ensuring that the network satisfies domain-specific
+requirements during both training and evaluation.
+
+The `CongradsCore` class serves as the central engine for managing the
+training, validation, and testing phases of a neural network model,
+incorporating constraints that influence the loss function and model updates.
+The model is trained with standard loss functions while also incorporating
+constraint-based adjustments, which are tracked and logged
+throughout the process.
+
+Key features:
+- Support for various constraints that can influence the training process.
+- Integration with PyTorch's `DataLoader` for efficient batch processing.
+- Metric management for tracking loss and constraint satisfaction.
+- Checkpoint management for saving and evaluating model states.
+
+"""
+
+from collections.abc import Callable
+
+import torch
+from torch import Tensor, sum
+from torch.nn import Module
+from torch.nn.modules.loss import _Loss
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from congrads.utils.utility import LossWrapper
+
+from ..callbacks.base import CallbackManager
+from ..checkpoints import CheckpointManager
+from ..constraints.base import Constraint
+from ..core.batch_runner import BatchRunner
+from ..core.constraint_engine import ConstraintEngine
+from ..core.epoch_runner import EpochRunner
+from ..descriptor import Descriptor
+from ..metrics import MetricManager
+
+
+class CongradsCore:
+    """The CongradsCore class is the central training engine for constraint-guided optimization.
+
+    It integrates standard neural network training
+    with additional constraint-driven adjustments to the loss function, ensuring
+    that the network satisfies domain-specific constraints during training.
+    """
+
+    def __init__(
+        self,
+        descriptor: Descriptor,
+        constraints: list[Constraint],
+        network: Module,
+        criterion: _Loss,
+        optimizer: Optimizer,
+        device: torch.device,
+        dataloader_train: DataLoader,
+        dataloader_valid: DataLoader | None = None,
+        dataloader_test: DataLoader | None = None,
+        metric_manager: MetricManager | None = None,
+        callback_manager: CallbackManager | None = None,
+        checkpoint_manager: CheckpointManager | None = None,
+        network_uses_grad: bool = False,
+        epsilon: float = 1e-6,
+        constraint_aggregator: Callable[..., Tensor] = sum,
+        enforce_all: bool = True,
+        disable_progress_bar_epoch: bool = False,
+        disable_progress_bar_batch: bool = False,
+        epoch_runner_cls: type["EpochRunner"] | None = None,
+        batch_runner_cls: type["BatchRunner"] | None = None,
+        constraint_engine_cls: type["ConstraintEngine"] | None = None,
+    ):
+        """Initialize the CongradsCore object.
+
+        Args:
+            descriptor (Descriptor): Describes variable layers in the network.
+            constraints (list[Constraint]): List of constraints to guide training.
+            network (Module): The neural network model to train.
+            criterion (callable): The loss function used for
+                training and validation.
+            optimizer (Optimizer): The optimizer used for updating model parameters.
+            device (torch.device): The device (e.g., CPU or GPU) for computations.
+            dataloader_train (DataLoader): DataLoader for training data.
+            dataloader_valid (DataLoader, optional): DataLoader for validation data.
+                If not provided, validation is skipped.
+            dataloader_test (DataLoader, optional): DataLoader for test data.
+                If not provided, testing is skipped.
+            metric_manager (MetricManager, optional): Manages metric tracking and recording.
+            callback_manager (CallbackManager, optional): Manages training callbacks.
+            checkpoint_manager (CheckpointManager, optional): Manages
+                    checkpointing. If not set, no checkpointing is done.
+            network_uses_grad (bool, optional): A flag indicating if the network
+                contains gradient calculation computations. Default is False.
+            epsilon (float, optional): A small value to avoid division by zero
+                in gradient calculations. Default is 1e-10.
+            constraint_aggregator (Callable[..., Tensor], optional): A function
+                to aggregate the constraint rescale loss. Default is `sum`.
+            enforce_all (bool, optional): If set to False, constraints will only be monitored and
+                not influence the training process. Overrides constraint-specific `enforce` parameters.
+                Defaults to True.
+            disable_progress_bar_epoch (bool, optional): If set to True, the epoch
+                progress bar will not show. Defaults to False.
+            disable_progress_bar_batch (bool, optional): If set to True, the batch
+                progress bar will not show. Defaults to False.
+            epoch_runner_cls (type[EpochRunner], optional): Custom EpochRunner class.
+                If not provided, the default EpochRunner is used.
+            batch_runner_cls (type[BatchRunner], optional): Custom BatchRunner class.
+                If not provided, the default BatchRunner is used.
+            constraint_engine_cls (type[ConstraintEngine], optional): Custom ConstraintEngine class.
+                If not provided, the default ConstraintEngine is used.
+
+        Note:
+            A warning is logged if the descriptor has no variable layers,
+            as at least one variable layer is required for the constraint logic
+            to influence the training process.
+        """
+        # Init object variables
+        self.device = device
+        self.network = network.to(device)
+        self.criterion = LossWrapper(criterion)
+        self.optimizer = optimizer
+        self.descriptor = descriptor
+
+        self.constraints = constraints or []
+        self.epsilon = epsilon
+        self.aggregator = constraint_aggregator
+        self.enforce_all = enforce_all
+
+        self.metric_manager = metric_manager
+        self.callback_manager = callback_manager
+        self.checkpoint_manager = checkpoint_manager
+
+        self.disable_progress_bar_epoch = disable_progress_bar_epoch
+        self.disable_progress_bar_batch = disable_progress_bar_batch
+
+        # Initialize constraint engine
+        self.constraint_engine = (constraint_engine_cls or ConstraintEngine)(
+            constraints=self.constraints,
+            descriptor=self.descriptor,
+            device=self.device,
+            epsilon=self.epsilon,
+            aggregator=self.aggregator,
+            enforce_all=self.enforce_all,
+            metric_manager=self.metric_manager,
+        )
+
+        # Initialize runners
+        self.batch_runner = (batch_runner_cls or BatchRunner)(
+            network=self.network,
+            criterion=self.criterion,
+            optimizer=self.optimizer,
+            constraint_engine=self.constraint_engine,
+            metric_manager=self.metric_manager,
+            callback_manager=self.callback_manager,
+            device=self.device,
+        )
+
+        self.epoch_runner = (epoch_runner_cls or EpochRunner)(
+            batch_runner=self.batch_runner,
+            network=self.network,
+            train_loader=dataloader_train,
+            valid_loader=dataloader_valid,
+            test_loader=dataloader_test,
+            network_uses_grad=network_uses_grad,
+            disable_progress_bar=self.disable_progress_bar_batch,
+        )
+
+        # Initialize constraint metrics
+        if self.metric_manager is not None:
+            self._initialize_metrics()
+
+    def _initialize_metrics(self) -> None:
+        """Register metrics for loss, constraint satisfaction ratio (CSR), and constraints.
+
+        This method registers the following metrics:
+
+        - Loss/train: Training loss.
+        - Loss/valid: Validation loss.
+        - Loss/test: Test loss after training.
+        - CSR/train: Constraint satisfaction ratio during training.
+        - CSR/valid: Constraint satisfaction ratio during validation.
+        - CSR/test: Constraint satisfaction ratio after training.
+        - One metric per constraint, for both training and validation.
+
+        """
+        self.metric_manager.register("Loss/train", "during_training")
+        self.metric_manager.register("Loss/valid", "during_training")
+        self.metric_manager.register("Loss/test", "after_training")
+
+        if len(self.constraints) > 0:
+            self.metric_manager.register("CSR/train", "during_training")
+            self.metric_manager.register("CSR/valid", "during_training")
+            self.metric_manager.register("CSR/test", "after_training")
+
+        for constraint in self.constraints:
+            self.metric_manager.register(f"{constraint.name}/train", "during_training")
+            self.metric_manager.register(f"{constraint.name}/valid", "during_training")
+            self.metric_manager.register(f"{constraint.name}/test", "after_training")
+
+    def fit(
+        self,
+        *,
+        start_epoch: int = 0,
+        max_epochs: int = 100,
+        test_model: bool = True,
+        final_checkpoint_name: str = "checkpoint_final.pth",
+    ) -> None:
+        """Run the full training loop, including optional validation, testing, and checkpointing.
+
+        This method performs training over multiple epochs with the following steps:
+        1. Trigger "on_train_start" callbacks if a callback manager is present.
+        2. For each epoch:
+        - Trigger "on_epoch_start" callbacks.
+        - Run a training epoch via the EpochRunner.
+        - Run a validation epoch via the EpochRunner.
+        - Evaluate checkpoint criteria if a checkpoint manager is present.
+        - Trigger "on_epoch_end" callbacks.
+        3. Trigger "on_train_end" callbacks after all epochs.
+        4. Optionally run a test epoch via the EpochRunner if `test_model` is True,
+        with corresponding "on_test_start" and "on_test_end" callbacks.
+        5. Save a final checkpoint using the checkpoint manager.
+
+        Args:
+            start_epoch: Index of the starting epoch (default 0). Useful for resuming training.
+            max_epochs: Maximum number of epochs to run (default 100).
+            test_model: Whether to run a test epoch after training (default True).
+            final_checkpoint_name: Filename for the final checkpoint saved at the end of training
+                                (default "checkpoint_final.pth").
+
+        Returns:
+            None
+        """
+        if self.callback_manager:
+            self.callback_manager.run("on_train_start", {"epoch": start_epoch})
+
+        for epoch in tqdm(
+            range(start_epoch, max_epochs),
+            initial=start_epoch,
+            desc="Epoch",
+            disable=self.disable_progress_bar_epoch,
+        ):
+            if self.callback_manager:
+                self.callback_manager.run("on_epoch_start", {"epoch": epoch})
+
+            self.epoch_runner.train()
+            self.epoch_runner.validate()
+
+            if self.checkpoint_manager:
+                self.checkpoint_manager.evaluate_criteria(epoch)
+
+            if self.callback_manager:
+                self.callback_manager.run("on_epoch_end", {"epoch": epoch})
+
+        if self.callback_manager:
+            self.callback_manager.run("on_train_end", {"epoch": epoch})
+
+        if test_model:
+            if self.callback_manager:
+                self.callback_manager.run("on_test_start", {"epoch": epoch})
+
+            self.epoch_runner.test()
+
+            if self.callback_manager:
+                self.callback_manager.run("on_test_end", {"epoch": epoch})
+
+        if self.checkpoint_manager:
+            self.checkpoint_manager.save(epoch, final_checkpoint_name)

congrads/core/constraint_engine.py

@@ -0,0 +1,209 @@
+"""Manages the evaluation and optional enforcement of constraints on neural network outputs.
+
+Responsibilities:
+- Compute and log Constraint Satisfaction Rate (CSR) for training, validation, and test batches.
+- Optionally adjust loss during training based on constraint directions and rescale factors.
+- Handle gradient computation and CGGD application.
+"""
+
+import torch
+from torch import Tensor, no_grad
+from torch.linalg import vector_norm
+
+from ..constraints.base import Constraint
+from ..descriptor import Descriptor
+from ..metrics import MetricManager
+
+
+class ConstraintEngine:
+    """Manages constraint evaluation and enforcement for a neural network.
+
+    The ConstraintEngine coordinates constraints defined in Constraint objects,
+    computes gradients for layers that affect the loss, logs metrics, and optionally
+    modifies the loss during training according to the constraints. It supports
+    separate phases for training, validation, and testing.
+    """
+
+    def __init__(
+        self,
+        *,
+        constraints: list[Constraint],
+        descriptor: Descriptor,
+        metric_manager: MetricManager,
+        device: torch.device,
+        epsilon: float,
+        aggregator: callable,
+        enforce_all: bool,
+    ) -> None:
+        """Initialize the ConstraintEngine.
+
+        Args:
+            constraints: List of Constraint objects to evaluate and optionally enforce.
+            descriptor: Descriptor containing metadata about network layers and which
+                        variables affect the loss.
+            metric_manager: MetricManager instance for logging CSR metrics.
+            device: Torch device where tensors will be allocated (CPU or GPU).
+            epsilon: Small positive value to avoid division by zero in gradient norms.
+            aggregator: Callable used to reduce per-layer constraint contributions
+                        to a scalar loss adjustment.
+            enforce_all: Whether to enforce all constraints during training.
+        """
+        self.constraints = constraints
+        self.descriptor = descriptor
+        self.metric_manager = metric_manager
+        self.device = device
+        self.epsilon = epsilon
+        self.enforce_all = enforce_all
+        self.aggregator = aggregator
+
+        self.norm_loss_grad: dict[str, Tensor] = {}
+
+    def train(self, data: dict[str, Tensor], loss: Tensor) -> Tensor:
+        """Apply all active constraints during training.
+
+        Computes the original loss gradients for layers that affect the loss,
+        evaluates each constraint, logs the Constraint Satisfaction Rate (CSR),
+        and adjusts the loss according to constraint satisfaction.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: The original loss tensor computed from the network output.
+
+        Returns:
+            Tensor: The loss tensor after applying constraint-based adjustments.
+        """
+        return self._apply_constraints(data, loss, phase="train", enforce=True)
+
+    def validate(self, data: dict[str, Tensor], loss: Tensor) -> Tensor:
+        """Evaluate constraints during validation without modifying the loss.
+
+        Computes and logs the Constraint Satisfaction Rate (CSR) for each constraint,
+        but does not apply rescale adjustments to the loss.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: The original loss tensor computed from the network output.
+
+        Returns:
+            Tensor: The original loss tensor, unchanged.
+        """
+        return self._apply_constraints(data, loss, phase="valid", enforce=False)
+
+    def test(self, data: dict[str, Tensor], loss: Tensor) -> Tensor:
+        """Evaluate constraints during testing without modifying the loss.
+
+        Computes and logs the Constraint Satisfaction Rate (CSR) for each constraint,
+        but does not apply rescale adjustments to the loss.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: The original loss tensor computed from the network output.
+
+        Returns:
+            Tensor: The original loss tensor, unchanged.
+        """
+        return self._apply_constraints(data, loss, phase="test", enforce=False)
+
+    def _apply_constraints(
+        self, data: dict[str, Tensor], loss: Tensor, phase: str, enforce: bool
+    ) -> Tensor:
+        """Evaluate constraints, log CSR, and optionally adjust the loss.
+
+        During training, computes loss gradients for variable layers that affect the loss.
+        Iterates over all constraints, logging the Constraint Satisfaction Rate (CSR)
+        and, if enforcement is enabled, adjusts the loss using constraint-specific
+        directions and rescale factors.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: Original loss tensor computed from the network output.
+            phase: Current phase, one of "train", "valid", or "test".
+            enforce: If True, constraint-based adjustments are applied to the loss.
+
+        Returns:
+            Tensor: The combined loss after applying constraints (or the original loss
+            if enforce is False or not in training phase).
+        """
+        total_rescale_loss = torch.tensor(0.0, device=self.device, dtype=loss.dtype)
+
+        if phase == "train":
+            norm_loss_grad = self._calculate_loss_gradients(loss, data)
+            norm_loss_grad = self._override_loss_gradients(norm_loss_grad, loss, data)
+
+        # Iterate constraints
+        for constraint in self.constraints:
+            checks, mask = constraint.check_constraint(data)
+            directions = constraint.calculate_direction(data)
+
+            # Log CSR
+            csr = (torch.sum(checks * mask) / torch.sum(mask)).unsqueeze(0)
+            self.metric_manager.accumulate(f"{constraint.name}/{phase}", csr)
+            self.metric_manager.accumulate(f"CSR/{phase}", csr)
+
+            # Skip adjustment if not enforcing
+            if not enforce or not constraint.enforce or not self.enforce_all or phase != "train":
+                continue
+
+            # Compute constraint-based rescale loss
+            for key in constraint.layers & self.descriptor.variable_keys:
+                with no_grad():
+                    rescale = (1 - checks) * directions[key] * constraint.rescale_factor
+                total_rescale_loss += self.aggregator(data[key] * rescale * norm_loss_grad[key])
+
+        return loss + total_rescale_loss
+
+    def _calculate_loss_gradients(self, loss: Tensor, data: dict[str, Tensor]) -> None:
+        """Compute and store normalized loss gradients for variable layers.
+
+        For each layer that affects the loss, computes the gradient of the loss
+        with respect to that layer's output. The gradients are normalized by their
+        vector norms plus a small epsilon to avoid division by zero.
+
+        Args:
+            loss: The original loss tensor computed from the network output.
+            data: Dictionary containing input and prediction tensors for the batch.
+        """
+        # Precompute gradients for variable layers affecting the loss
+        norm_loss_grad = {}
+
+        variable_keys = self.descriptor.variable_keys & self.descriptor.affects_loss_keys
+        for key in variable_keys:
+            if data[key].requires_grad is False:
+                raise RuntimeError(
+                    f"Layer '{key}' does not require gradients. Is this an input? "
+                    "Set constant=True in Descriptor if this layer is an input."
+                )
+
+            grad = torch.autograd.grad(
+                outputs=loss, inputs=data[key], retain_graph=True, allow_unused=True
+            )[0]
+
+            if grad is None:
+                raise RuntimeError(
+                    f"Unable to compute loss gradients for layer '{key}'. "
+                    "Set has_loss=False in Descriptor if this layer does not affect loss."
+                )
+
+            grad_flat = grad.view(grad.shape[0], -1)
+            norm_loss_grad[key] = (
+                vector_norm(grad_flat, dim=1, ord=2, keepdim=True).clamp(min=self.epsilon).detach()
+            )
+
+        return norm_loss_grad
+
+    def _override_loss_gradients(
+        self, norm_loss_grad: dict[str, Tensor], loss: Tensor, data: dict[str, Tensor]
+    ) -> dict[str, Tensor]:
+        """Override the standard normalized loss gradient computation for custom functionality.
+
+        Args:
+            norm_loss_grad: Dictionary mapping parameter or component names to normalized
+                gradient tensors.
+            loss: Scalar loss value for the current training step.
+            data: Dictionary containing the batch data used to compute the loss and any
+                constraint-related signals.
+
+        Returns:
+            A dictionary of modified normalized gradients.
+        """
+        return norm_loss_grad