Class TiedLinear

astroclip/modules.py:296–338

A dense linear layer whose parameters are tied to a tensor provided by the user. Using this layer is equivalent to using the functional form, :func:`~torch.nn.functional.linear`. The utility of having a module is that it will show up in module summaries, which can help to make the struc

class TiedLinear(nn.Module):
    """A dense linear layer whose parameters are tied to a tensor provided by the user.

    Using this layer is equivalent to using the functional form,
    :func:`~torch.nn.functional.linear`. The utility of having a module is that it will
    show up in module summaries, which can help to make the structure of the model more
    transparent.

    :param weight: weight tensor
    :param bias: bias tensor; if not provided, there will be no bias
    """

    in_features: int
    """size of each input sample."""

    out_features: int
    """size of each output sample."""

    def __init__(
        self,
        weight: Union[torch.Tensor, nn.Parameter],
        bias: Union[None, torch.Tensor, nn.Parameter],
    ):
        super().__init__()

        if weight.ndim != 2:
            raise ValueError(
                f"weight parameter has {weight.ndim} dimensions, should have 2"
            )
        self.out_features, self.in_features = weight.shape

        self.register_buffer("weight", weight)
        self.register_buffer("bias", bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.linear(x, self.weight, self.bias)

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, "
            f"out_features={self.out_features}, "
            f"bias={self.bias is not None}"
        )


def _init_by_depth(module: nn.Module, depth: int) -> None: