logit_layer.py

"""HF-local logit rescaling for Hub-compatible ALOE image-classification models.

Self-contained — no dependency on project ``src.*`` or the ``bcos`` PyPI package, so the
file stays loadable after it is copied into the shared Hub code repo. ``src.modules.logit_layer``
re-exports this class, keeping training and Hub inference on one definition.
"""

from __future__ import annotations

from typing import Optional

import torch
import torch.nn as nn


class LogitLayer(nn.Module):
    """
    Applies optional temperature scaling and bias to logits.
    Commonly used in B-cos models to adjust the dynamic range of logits.

    Parameter-free by design: ``logit_temperature``/``logit_bias`` are plain Python floats, so
    the module contributes no ``state_dict`` entries and checkpoints hold no ``logit_layer.*``
    weights. Both values come from the config at construction time.
    """

    def __init__(self, logit_temperature: Optional[float] = None, logit_bias: Optional[float] = None):
        super().__init__()
        self.logit_temperature = logit_temperature
        self.logit_bias = logit_bias

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.logit_temperature is not None:
            x = x * self.logit_temperature
        if self.logit_bias is not None:
            x = x + self.logit_bias
        return x