

"""simple docstring"""




from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)







class      __lowercase			(    __lowercase					):

  '''simple docstring'''





  def __init__(  self		:			List[Any]					,A		:			Distribution					,A		:			List[Any]=None					,A		:			Optional[int]=None					,A		:			Union[str, Any]=0    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Any					   =					1.0 if scale is None else scale
   UpperCAmelCase__					:  int					   =					0.0 if loc is None else loc

   super().__init__(__a					,[AffineTransform(loc=self.loc					,scale=self.scale					,event_dim=__a    )]    )

  @property
  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   return self.base_dist.mean * self.scale + self.loc

  @property
  def    __lowercase       (  self		:			Dict    ):






   '''simple docstring'''
   return self.base_dist.variance * self.scale**2




  @property
  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   return self.variance.sqrt()







class      __lowercase			(    nn.Module					):

  '''simple docstring'''





  def __init__(  self		:			Optional[int]					,A		:			int					,A		:			Dict[str, int]					,A		:			Callable[..., Tuple[torch.Tensor]]					,**A		:			Optional[Any]    ):






   '''simple docstring'''
   super().__init__(**__a    )
   UpperCAmelCase__					:  Optional[int]					   =					args_dim
   UpperCAmelCase__					:  List[str]					   =					nn.ModuleList([nn.Linear(__a					,__a    ) for dim in args_dim.values()]    )
   UpperCAmelCase__					:  str					   =					domain_map




  def    __lowercase       (  self		:			List[str]					,A		:			torch.Tensor    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					[proj(__a    ) for proj in self.proj]

   return self.domain_map(*__a    )







class      __lowercase			(    nn.Module					):

  '''simple docstring'''





  def __init__(  self		:			Union[str, Any]					,A		:			str    ):






   '''simple docstring'''
   super().__init__()
   UpperCAmelCase__					:  Any					   =					function




  def    __lowercase       (  self		:			str					,A		:			Union[str, Any]					,*A		:			Optional[Any]    ):






   '''simple docstring'''
   return self.function(__a					,*__a    )







class      __lowercase			:

  '''simple docstring'''





  snake_case_				       =       42
  snake_case_				       =       42
  snake_case_				       =       42

  def __init__(  self		:			Dict					,A		:			int = 1    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					dim
   UpperCAmelCase__					:  List[Any]					   =					{k: dim * self.args_dim[k] for k in self.args_dim}

  def    __lowercase       (  self		:			List[str]					,A		:			Dict    ):






   '''simple docstring'''
   if self.dim == 1:
    return self.distribution_class(*__a    )
   else:
    return Independent(self.distribution_class(*__a    )					,1    )

  def    __lowercase       (  self		:			Any					,A		:			str					,A		:			Optional[torch.Tensor] = None					,A		:			Optional[torch.Tensor] = None					,):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					self._base_distribution(__a    )
   if loc is None and scale is None:
    return distr
   else:
    return AffineTransformed(__a					,loc=__a					,scale=__a					,event_dim=self.event_dim    )

  @property
  def    __lowercase       (  self		:			str    ):






   '''simple docstring'''
   return () if self.dim == 1 else (self.dim,)

  @property
  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   return len(self.event_shape    )

  @property
  def    __lowercase       (  self		:			Dict    ):






   '''simple docstring'''
   return 0.0

  def    __lowercase       (  self		:			Union[str, Any]					,A		:			int    ):






   '''simple docstring'''
   return ParameterProjection(
       in_features=__a					,args_dim=self.args_dim					,domain_map=LambdaLayer(self.domain_map    )					,)

  def    __lowercase       (  self		:			Any					,*A		:			torch.Tensor    ):






   '''simple docstring'''
   raise NotImplementedError()




  @staticmethod
  def    __lowercase       (  A		:			torch.Tensor    ):






   '''simple docstring'''
   return (x + torch.sqrt(torch.square(__a    ) + 4.0    )) / 2.0







class      __lowercase			(    __lowercase					):

  '''simple docstring'''





  snake_case_				       =       {"df": 1, "loc": 1, "scale": 1}
  snake_case_				       =       StudentT

  @classmethod
  def    __lowercase       (  cls		:			Optional[Any]					,A		:			torch.Tensor					,A		:			torch.Tensor					,A		:			torch.Tensor    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[Any]					   =					cls.squareplus(__a    ).clamp_min(torch.finfo(scale.dtype    ).eps    )
   UpperCAmelCase__					:  Optional[int]					   =					2.0 + cls.squareplus(__a    )
   return df.squeeze(-1    ), loc.squeeze(-1    ), scale.squeeze(-1    )







class      __lowercase			(    __lowercase					):

  '''simple docstring'''





  snake_case_				       =       {"loc": 1, "scale": 1}
  snake_case_				       =       Normal

  @classmethod
  def    __lowercase       (  cls		:			Optional[int]					,A		:			torch.Tensor					,A		:			torch.Tensor    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					cls.squareplus(__a    ).clamp_min(torch.finfo(scale.dtype    ).eps    )
   return loc.squeeze(-1    ), scale.squeeze(-1    )







class      __lowercase			(    __lowercase					):

  '''simple docstring'''





  snake_case_				       =       {"total_count": 1, "logits": 1}
  snake_case_				       =       NegativeBinomial

  @classmethod
  def    __lowercase       (  cls		:			Tuple					,A		:			torch.Tensor					,A		:			torch.Tensor    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					cls.squareplus(__a    )
   return total_count.squeeze(-1    ), logits.squeeze(-1    )

  def    __lowercase       (  self		:			List[Any]					,A		:			Union[str, Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					distr_args
   if self.dim == 1:
    return self.distribution_class(total_count=__a					,logits=__a    )
   else:
    return Independent(self.distribution_class(total_count=__a					,logits=__a    )					,1    )




  def    __lowercase       (  self		:			str					,A		:			Any					,A		:			Optional[torch.Tensor] = None					,A		:			Optional[torch.Tensor] = None    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[str]					   =					distr_args

   if scale is not None:
    # See scaling property of Gamma.
    logits += scale.log()

   return self._base_distribution((total_count, logits)    )








"""simple docstring"""




def   lowerCAmelCase						(					__UpperCamelCase = 1000      ):
 '''simple docstring'''





 return sum(2 * a * ((a - 1) // 2) for a in range(3							,    n + 1      )      )


if __name__ == "__main__":
      print(solution())





