Aluode/PerceptionLabPortable

1# Copyright 2023 The HuggingFace Inc. team.2# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""16Time series distributional output classes and utilities.17"""18 19from typing import Callable, Optional20 21import torch22from torch import nn23from torch.distributions import (24    AffineTransform,25    Distribution,26    Independent,27    NegativeBinomial,28    Normal,29    StudentT,30    TransformedDistribution,31)32 33 34class AffineTransformed(TransformedDistribution):35    def __init__(self, base_distribution: Distribution, loc=None, scale=None, event_dim=0):36        self.scale = 1.0 if scale is None else scale37        self.loc = 0.0 if loc is None else loc38 39        super().__init__(base_distribution, [AffineTransform(loc=self.loc, scale=self.scale, event_dim=event_dim)])40 41    @property42    def mean(self):43        """44        Returns the mean of the distribution.45        """46        return self.base_dist.mean * self.scale + self.loc47 48    @property49    def variance(self):50        """51        Returns the variance of the distribution.52        """53        return self.base_dist.variance * self.scale**254 55    @property56    def stddev(self):57        """58        Returns the standard deviation of the distribution.59        """60        return self.variance.sqrt()61 62 63class ParameterProjection(nn.Module):64    def __init__(65        self, in_features: int, args_dim: dict[str, int], domain_map: Callable[..., tuple[torch.Tensor]], **kwargs66    ) -> None:67        super().__init__(**kwargs)68        self.args_dim = args_dim69        self.proj = nn.ModuleList([nn.Linear(in_features, dim) for dim in args_dim.values()])70        self.domain_map = domain_map71 72    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor]:73        params_unbounded = [proj(x) for proj in self.proj]74 75        return self.domain_map(*params_unbounded)76 77 78class LambdaLayer(nn.Module):79    def __init__(self, function):80        super().__init__()81        self.function = function82 83    def forward(self, x, *args):84        return self.function(x, *args)85 86 87class DistributionOutput:88    distribution_class: type89    in_features: int90    args_dim: dict[str, int]91 92    def __init__(self, dim: int = 1) -> None:93        self.dim = dim94        self.args_dim = {k: dim * self.args_dim[k] for k in self.args_dim}95 96    def _base_distribution(self, distr_args):97        if self.dim == 1:98            return self.distribution_class(*distr_args)99        else:100            return Independent(self.distribution_class(*distr_args), 1)101 102    def distribution(103        self,104        distr_args,105        loc: Optional[torch.Tensor] = None,106        scale: Optional[torch.Tensor] = None,107    ) -> Distribution:108        distr = self._base_distribution(distr_args)109        if loc is None and scale is None:110            return distr111        else:112            return AffineTransformed(distr, loc=loc, scale=scale, event_dim=self.event_dim)113 114    @property115    def event_shape(self) -> tuple:116        r"""117        Shape of each individual event contemplated by the distributions that this object constructs.118        """119        return () if self.dim == 1 else (self.dim,)120 121    @property122    def event_dim(self) -> int:123        r"""124        Number of event dimensions, i.e., length of the `event_shape` tuple, of the distributions that this object125        constructs.126        """127        return len(self.event_shape)128 129    @property130    def value_in_support(self) -> float:131        r"""132        A float that will have a valid numeric value when computing the log-loss of the corresponding distribution. By133        default 0.0. This value will be used when padding data series.134        """135        return 0.0136 137    def get_parameter_projection(self, in_features: int) -> nn.Module:138        r"""139        Return the parameter projection layer that maps the input to the appropriate parameters of the distribution.140        """141        return ParameterProjection(142            in_features=in_features,143            args_dim=self.args_dim,144            domain_map=LambdaLayer(self.domain_map),145        )146 147    def domain_map(self, *args: torch.Tensor):148        r"""149        Converts arguments to the right shape and domain. The domain depends on the type of distribution, while the150        correct shape is obtained by reshaping the trailing axis in such a way that the returned tensors define a151        distribution of the right event_shape.152        """153        raise NotImplementedError()154 155    @staticmethod156    def squareplus(x: torch.Tensor) -> torch.Tensor:157        r"""158        Helper to map inputs to the positive orthant by applying the square-plus operation. Reference:159        https://twitter.com/jon_barron/status/1387167648669048833160        """161        return (x + torch.sqrt(torch.square(x) + 4.0)) / 2.0162 163 164class StudentTOutput(DistributionOutput):165    """166    Student-T distribution output class.167    """168 169    args_dim: dict[str, int] = {"df": 1, "loc": 1, "scale": 1}170    distribution_class: type = StudentT171 172    @classmethod173    def domain_map(cls, df: torch.Tensor, loc: torch.Tensor, scale: torch.Tensor):174        scale = cls.squareplus(scale).clamp_min(torch.finfo(scale.dtype).eps)175        df = 2.0 + cls.squareplus(df)176        return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1)177 178 179class NormalOutput(DistributionOutput):180    """181    Normal distribution output class.182    """183 184    args_dim: dict[str, int] = {"loc": 1, "scale": 1}185    distribution_class: type = Normal186 187    @classmethod188    def domain_map(cls, loc: torch.Tensor, scale: torch.Tensor):189        scale = cls.squareplus(scale).clamp_min(torch.finfo(scale.dtype).eps)190        return loc.squeeze(-1), scale.squeeze(-1)191 192 193class NegativeBinomialOutput(DistributionOutput):194    """195    Negative Binomial distribution output class.196    """197 198    args_dim: dict[str, int] = {"total_count": 1, "logits": 1}199    distribution_class: type = NegativeBinomial200 201    @classmethod202    def domain_map(cls, total_count: torch.Tensor, logits: torch.Tensor):203        total_count = cls.squareplus(total_count)204        return total_count.squeeze(-1), logits.squeeze(-1)205 206    def _base_distribution(self, distr_args) -> Distribution:207        total_count, logits = distr_args208        if self.dim == 1:209            return self.distribution_class(total_count=total_count, logits=logits)210        else:211            return Independent(self.distribution_class(total_count=total_count, logits=logits), 1)212 213    # Overwrites the parent class method. We cannot scale using the affine214    # transformation since negative binomial should return integers. Instead215    # we scale the parameters.216    def distribution(217        self, distr_args, loc: Optional[torch.Tensor] = None, scale: Optional[torch.Tensor] = None218    ) -> Distribution:219        total_count, logits = distr_args220 221        if scale is not None:222            # See scaling property of Gamma.223            logits += scale.log()224 225        return self._base_distribution((total_count, logits))226