






'''simple docstring'''







import math
from typing import Optional

import numpy as np

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


UpperCamelCase     :   List[Any]	      =		logging.get_logger(__name__)

UpperCamelCase     :   str	      =		{
    '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__     (a      ):




		'''simple docstring'''





		_UpperCamelCase						=       'encodec'






		def __init__(   self ,_lowerCAmelCase=[1.5, 3.0, 6.0, 12.0, 24.0] ,_lowerCAmelCase=2_40_00 ,_lowerCAmelCase=1 ,_lowerCAmelCase=False ,_lowerCAmelCase=None ,_lowerCAmelCase=None ,_lowerCAmelCase=1_28 ,_lowerCAmelCase=32 ,_lowerCAmelCase=1 ,_lowerCAmelCase=[8, 5, 4, 2] ,_lowerCAmelCase="weight_norm" ,_lowerCAmelCase=7 ,_lowerCAmelCase=7 ,_lowerCAmelCase=3 ,_lowerCAmelCase=2 ,_lowerCAmelCase=True ,_lowerCAmelCase="reflect" ,_lowerCAmelCase=2 ,_lowerCAmelCase=2 ,_lowerCAmelCase=1.0 ,_lowerCAmelCase=10_24 ,_lowerCAmelCase=None ,_lowerCAmelCase=True ,**_lowerCAmelCase ,):
								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       UpperCamelCase_       (   self							):
								if self.chunk_length_s is None:
														return None
								else:
														return int(self.chunk_length_s * self.sampling_rate							)






		@property
		def       UpperCamelCase_       (   self							):
								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       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	np.prod(self.upsampling_ratios							)
								return math.ceil(self.sampling_rate / hop_length							)






		@property
		def       UpperCamelCase_       (   self							):
								return int(10_00 * self.target_bandwidths[-1] // (self.frame_rate * 10)							)









'''simple docstring'''







import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__lowerCAmelCase									=   object()

# For specifying empty leaf dict `{}`
__lowerCAmelCase									=   object()





def 		_UpperCAmelCase (				__A				: List[str]    ,		__A				: Tuple   ):
				a_		:	List[Any]									=   tuple((re.compile(x + '''$'''   ) for x in qs)   )
				for i in range(len(__A   ) - len(__A   ) + 1   ):
								a_		:	Union[str, Any]									=   [x.match(__A   ) for x, y in zip(__A    ,		ks[i:]   )]
								if matches and all(__A   ):
												return True
				return False





def 		_UpperCAmelCase (				__A				: List[str]   ):
				def replace(__A				: int    ,		__A				: Union[str, Any]   ):
								for rule, replacement in rules:
												if _match(__A    ,		__A   ):
																return replacement
								return val

				return replace





def 		_UpperCAmelCase (				):
				return [
				    # embeddings
				    (("transformer", "wpe", "embedding"), P('''mp'''    ,		__A   )),
				    (("transformer", "wte", "embedding"), P('''mp'''    ,		__A   )),
				    # atention
				    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__A    ,		'''mp'''   )),
				    (("attention", "out_proj", "kernel"), P('''mp'''    ,		__A   )),
				    (("attention", "out_proj", "bias"), None),
				    # mlp
				    (("mlp", "c_fc", "kernel"), P(__A    ,		'''mp'''   )),
				    (("mlp", "c_fc", "bias"), P('''mp'''   )),
				    (("mlp", "c_proj", "kernel"), P('''mp'''    ,		__A   )),
				    (("mlp", "c_proj", "bias"), None),
				    # layer norms
				    ((r"ln_\d+", "bias"), None),
				    ((r"\d+", r"ln_\d+", "scale"), None),
				    (("ln_f", "bias"), None),
				    (("ln_f", "scale"), None),
				]





def 		_UpperCAmelCase (				__A				: Union[str, Any]   ):
				a_		:	Tuple									=   _get_partition_rules()
				a_		:	Tuple									=   _replacement_rules(__A   )
				a_		:	Optional[Any]									=   {k: _unmatched for k in flatten_dict(__A   )}
				a_		:	Optional[Any]									=   {k: replace(__A    ,		__A   ) for k, v in initd.items()}
				assert _unmatched not in result.values(), "Incomplete partition spec."
				return freeze(unflatten_dict(__A   )   )




