





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     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  def __init__(    self			:		int ,  lowerCAmelCase_			:		Distribution ,  lowerCAmelCase_			:		Dict=None ,  lowerCAmelCase_			:		Tuple=None ,  lowerCAmelCase_			:		Union[str, Any]=0					)							->						str:
     __lowerCAmelCase							     =     1.0 if scale is None else scale
     __lowerCAmelCase							     =     0.0 if loc is None else loc

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






  @property
  def     lowercase		(    self			:		int					)							->						Dict:
     return self.base_dist.mean * self.scale + self.loc






  @property
  def     lowercase		(    self			:		Dict					)							->						int:
     return self.base_dist.variance * self.scale**2






  @property
  def     lowercase		(    self			:		int					)							->						Optional[int]:
     return self.variance.sqrt()





class     _UpperCAmelCase   (   nn.Module   ):

  """simple docstring"""


  def __init__(    self			:		str ,  lowerCAmelCase_			:		int ,  lowerCAmelCase_			:		Dict[str, int] ,  lowerCAmelCase_			:		Callable[..., Tuple[torch.Tensor]] ,  **lowerCAmelCase_			:		List[str]					)							->						None:
     super().__init__(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     args_dim
     __lowerCAmelCase							     =     nn.ModuleList([nn.Linear(lowerCAmelCase_ ,  lowerCAmelCase_					) for dim in args_dim.values()]					)
     __lowerCAmelCase							     =     domain_map






  def     lowercase		(    self			:		Dict ,  lowerCAmelCase_			:		torch.Tensor					)							->						Tuple[torch.Tensor]:
     __lowerCAmelCase							     =     [proj(lowerCAmelCase_					) for proj in self.proj]

     return self.domain_map(*lowerCAmelCase_					)





class     _UpperCAmelCase   (   nn.Module   ):

  """simple docstring"""


  def __init__(    self			:		Union[str, Any] ,  lowerCAmelCase_			:		List[str]					)							->						List[Any]:
     super().__init__()
     __lowerCAmelCase							     =     function






  def     lowercase		(    self			:		List[Any] ,  lowerCAmelCase_			:		int ,  *lowerCAmelCase_			:		List[Any]					)							->						List[Any]:
     return self.function(lowerCAmelCase_ ,  *lowerCAmelCase_					)





class     _UpperCAmelCase   :

  """simple docstring"""


  a_      =			42
  a_      =			42
  a_      =			42






  def __init__(    self			:		Dict ,  lowerCAmelCase_			:		int = 1					)							->						None:
     __lowerCAmelCase							     =     dim
     __lowerCAmelCase							     =     {k: dim * self.args_dim[k] for k in self.args_dim}






  def     lowercase		(    self			:		str ,  lowerCAmelCase_			:		Union[str, Any]					)							->						Union[str, Any]:
     if self.dim == 1:
        return self.distribution_class(*lowerCAmelCase_					)
     else:
        return Independent(self.distribution_class(*lowerCAmelCase_					) ,  1					)






  def     lowercase		(    self			:		Any ,  lowerCAmelCase_			:		Union[str, Any] ,  lowerCAmelCase_			:		Optional[torch.Tensor] = None ,  lowerCAmelCase_			:		Optional[torch.Tensor] = None ,  )							->						Distribution:
     __lowerCAmelCase							     =     self._base_distribution(lowerCAmelCase_					)
     if loc is None and scale is None:
        return distr
     else:
        return AffineTransformed(lowerCAmelCase_ ,  loc=lowerCAmelCase_ ,  scale=lowerCAmelCase_ ,  event_dim=self.event_dim					)






  @property
  def     lowercase		(    self			:		Optional[int]					)							->						Tuple:
     return () if self.dim == 1 else (self.dim,)






  @property
  def     lowercase		(    self			:		List[Any]					)							->						int:
     return len(self.event_shape					)






  @property
  def     lowercase		(    self			:		List[Any]					)							->						float:
     return 0.0






  def     lowercase		(    self			:		str ,  lowerCAmelCase_			:		int					)							->						nn.Module:
     return ParameterProjection(
         in_features=lowerCAmelCase_ ,  args_dim=self.args_dim ,  domain_map=LambdaLayer(self.domain_map					) ,  )






  def     lowercase		(    self			:		Optional[int] ,  *lowerCAmelCase_			:		torch.Tensor					)							->						int:
     raise NotImplementedError()






  @staticmethod
  def     lowercase		(    lowerCAmelCase_			:		torch.Tensor					)							->						torch.Tensor:
     return (x + torch.sqrt(torch.square(lowerCAmelCase_					) + 4.0					)) / 2.0





class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


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






  @classmethod
  def     lowercase		(    cls			:		List[Any] ,  lowerCAmelCase_			:		torch.Tensor ,  lowerCAmelCase_			:		torch.Tensor ,  lowerCAmelCase_			:		torch.Tensor					)							->						List[str]:
     __lowerCAmelCase							     =     cls.squareplus(lowerCAmelCase_					).clamp_min(torch.finfo(scale.dtype					).eps					)
     __lowerCAmelCase							     =     2.0 + cls.squareplus(lowerCAmelCase_					)
     return df.squeeze(-1					), loc.squeeze(-1					), scale.squeeze(-1					)





class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


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






  @classmethod
  def     lowercase		(    cls			:		str ,  lowerCAmelCase_			:		torch.Tensor ,  lowerCAmelCase_			:		torch.Tensor					)							->						str:
     __lowerCAmelCase							     =     cls.squareplus(lowerCAmelCase_					).clamp_min(torch.finfo(scale.dtype					).eps					)
     return loc.squeeze(-1					), scale.squeeze(-1					)





class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


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






  @classmethod
  def     lowercase		(    cls			:		List[Any] ,  lowerCAmelCase_			:		torch.Tensor ,  lowerCAmelCase_			:		torch.Tensor					)							->						Dict:
     __lowerCAmelCase							     =     cls.squareplus(lowerCAmelCase_					)
     return total_count.squeeze(-1					), logits.squeeze(-1					)






  def     lowercase		(    self			:		Tuple ,  lowerCAmelCase_			:		Optional[Any]					)							->						Distribution:
     __lowerCAmelCase     ,     __lowerCAmelCase							     =     distr_args
     if self.dim == 1:
        return self.distribution_class(total_count=lowerCAmelCase_ ,  logits=lowerCAmelCase_					)
     else:
        return Independent(self.distribution_class(total_count=lowerCAmelCase_ ,  logits=lowerCAmelCase_					) ,  1					)






  def     lowercase		(    self			:		str ,  lowerCAmelCase_			:		Any ,  lowerCAmelCase_			:		Optional[torch.Tensor] = None ,  lowerCAmelCase_			:		Optional[torch.Tensor] = None					)							->						Distribution:
     __lowerCAmelCase     ,     __lowerCAmelCase							     =     distr_args

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

     return self._base_distribution((total_count, logits)					)








import math
from typing import Optional

import numpy as np

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_snake_case  :   List[str]						   =  logging.get_logger(__name__)

_snake_case  :   Dict						   =  {
    'facebook/encodec_24khz': 'https://huggingface.co/facebook/encodec_24khz/resolve/main/config.json',
    'facebook/encodec_48khz': 'https://huggingface.co/facebook/encodec_48khz/resolve/main/config.json',
}



class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  a_      =			"""encodec"""






  def __init__(    self			:		Optional[Any] ,  lowerCAmelCase_			:		List[Any]=[1.5, 3.0, 6.0, 12.0, 24.0] ,  lowerCAmelCase_			:		Optional[Any]=2_4_0_0_0 ,  lowerCAmelCase_			:		int=1 ,  lowerCAmelCase_			:		Union[str, Any]=False ,  lowerCAmelCase_			:		Optional[Any]=None ,  lowerCAmelCase_			:		Union[str, Any]=None ,  lowerCAmelCase_			:		Optional[Any]=1_2_8 ,  lowerCAmelCase_			:		List[Any]=3_2 ,  lowerCAmelCase_			:		Optional[Any]=1 ,  lowerCAmelCase_			:		Optional[int]=[8, 5, 4, 2] ,  lowerCAmelCase_			:		Optional[Any]="weight_norm" ,  lowerCAmelCase_			:		Tuple=7 ,  lowerCAmelCase_			:		Optional[int]=7 ,  lowerCAmelCase_			:		Dict=3 ,  lowerCAmelCase_			:		Optional[Any]=2 ,  lowerCAmelCase_			:		List[Any]=True ,  lowerCAmelCase_			:		Dict="reflect" ,  lowerCAmelCase_			:		List[Any]=2 ,  lowerCAmelCase_			:		int=2 ,  lowerCAmelCase_			:		Dict=1.0 ,  lowerCAmelCase_			:		str=1_0_2_4 ,  lowerCAmelCase_			:		Dict=None ,  lowerCAmelCase_			:		List[Any]=True ,  **lowerCAmelCase_			:		Optional[Any] ,  )							->						Union[str, Any]:
     __lowerCAmelCase							     =     target_bandwidths
     __lowerCAmelCase							     =     sampling_rate
     __lowerCAmelCase							     =     audio_channels
     __lowerCAmelCase							     =     normalize
     __lowerCAmelCase							     =     chunk_length_s
     __lowerCAmelCase							     =     overlap
     __lowerCAmelCase							     =     hidden_size
     __lowerCAmelCase							     =     num_filters
     __lowerCAmelCase							     =     num_residual_layers
     __lowerCAmelCase							     =     upsampling_ratios
     __lowerCAmelCase							     =     norm_type
     __lowerCAmelCase							     =     kernel_size
     __lowerCAmelCase							     =     last_kernel_size
     __lowerCAmelCase							     =     residual_kernel_size
     __lowerCAmelCase							     =     dilation_growth_rate
     __lowerCAmelCase							     =     use_causal_conv
     __lowerCAmelCase							     =     pad_mode
     __lowerCAmelCase							     =     compress
     __lowerCAmelCase							     =     num_lstm_layers
     __lowerCAmelCase							     =     trim_right_ratio
     __lowerCAmelCase							     =     codebook_size
     __lowerCAmelCase							     =     codebook_dim if codebook_dim is not None else hidden_size
     __lowerCAmelCase							     =     use_conv_shortcut

     if self.norm_type not in ["weight_norm", "time_group_norm"]:
        raise ValueError(
            f"""self.norm_type must be one of `\"weight_norm\"`, `\"time_group_norm\"`), got {self.norm_type}"""					)

     super().__init__(**lowerCAmelCase_					)






  @property
  def     lowercase		(    self			:		Union[str, Any]					)							->						Optional[int]:
     if self.chunk_length_s is None:
        return None
     else:
        return int(self.chunk_length_s * self.sampling_rate					)






  @property
  def     lowercase		(    self			:		Optional[Any]					)							->						Optional[int]:
     if self.chunk_length_s is None or self.overlap is None:
        return None
     else:
        return max(1 ,  int((1.0 - self.overlap) * self.chunk_length					)					)






  @property
  def     lowercase		(    self			:		Any					)							->						int:
     __lowerCAmelCase							     =     np.prod(self.upsampling_ratios					)
     return math.ceil(self.sampling_rate / hop_length					)






  @property
  def     lowercase		(    self			:		str					)							->						int:
     return int(1_0_0_0 * self.target_bandwidths[-1] // (self.frame_rate * 1_0)					)

