






"""simple docstring"""







import functools
import operator

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


__a       :	Tuple             = logging.get_logger(__name__)

__a       :	List[str]             = {
    'asapp/sew-d-tiny-100k': 'https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json',
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}







class        _SCREAMING_SNAKE_CASE						(  __snake_case							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    ='sew-d'




							def __init__(   self:     Optional[Any]				,		__A:     List[Any]=32				,		__A:     Dict=768				,		__A:     Any=12				,		__A:     int=12				,		__A:     Union[str, Any]=3072				,		__A:     List[Any]=2				,		__A:     Any=512				,		__A:     Optional[Any]=256				,		__A:     Optional[int]=True				,		__A:     str=True				,		__A:     List[Any]=("p2c", "c2p")				,		__A:     Union[str, Any]="layer_norm"				,		__A:     Optional[int]="gelu_python"				,		__A:     Dict=0.1				,		__A:     str=0.1				,		__A:     Optional[Any]=0.1				,		__A:     Dict=0.0				,		__A:     Tuple=0.1				,		__A:     Dict=0.0_2				,		__A:     Dict=1e-7				,		__A:     str=1e-5				,		__A:     Dict="group"				,		__A:     int="gelu"				,		__A:     int=(64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512)				,		__A:     Any=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)				,		__A:     List[Any]=(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)				,		__A:     Tuple=False				,		__A:     Optional[Any]=128				,		__A:     str=16				,		__A:     Tuple=True				,		__A:     Tuple=0.0_5				,		__A:     Optional[int]=10				,		__A:     Union[str, Any]=2				,		__A:     str=0.0				,		__A:     Union[str, Any]=10				,		__A:     Optional[Any]=0				,		__A:     List[Any]="mean"				,		__A:     Any=False				,		__A:     List[str]=False				,		__A:     Dict=256				,		__A:     Optional[Any]=0				,		__A:     int=1				,		__A:     Any=2				,		**__A:     int				,		):

												'''simple docstring'''

												super().__init__(**__A				,		pad_token_id=__A				,		bos_token_id=__A				,		eos_token_id=__A    )
												a__          =  hidden_size
												a__          =  feat_extract_norm
												a__          =  feat_extract_activation
												a__          =  list(__A    )
												a__          =  list(__A    )
												a__          =  list(__A    )
												a__          =  conv_bias
												a__          =  num_conv_pos_embeddings
												a__          =  num_conv_pos_embedding_groups
												a__          =  len(self.conv_dim    )
												a__          =  num_hidden_layers
												a__          =  intermediate_size
												a__          =  squeeze_factor
												a__          =  max_position_embeddings
												a__          =  position_buckets
												a__          =  share_att_key
												a__          =  relative_attention
												a__          =  norm_rel_ebd
												a__          =  list(__A    )
												a__          =  hidden_act
												a__          =  num_attention_heads
												a__          =  hidden_dropout
												a__          =  attention_dropout
												a__          =  activation_dropout
												a__          =  feat_proj_dropout
												a__          =  final_dropout
												a__          =  layer_norm_eps
												a__          =  feature_layer_norm_eps
												a__          =  initializer_range
												a__          =  vocab_size

												if (
												    (len(self.conv_stride    ) != self.num_feat_extract_layers)
												    or (len(self.conv_kernel    ) != self.num_feat_extract_layers)
												    or (len(self.conv_dim    ) != self.num_feat_extract_layers)
												):
																	raise ValueError(
																	    '''Configuration for convolutional layers is incorrect.'''
																	    '''It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'''
																	    F'but is `len(config.conv_dim) = {len(self.conv_dim    )}`, `len(config.conv_stride)'
																	    F'= {len(self.conv_stride    )}`, `len(config.conv_kernel) = {len(self.conv_kernel    )}`.'    )

												# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
												a__          =  apply_spec_augment
												a__          =  mask_time_prob
												a__          =  mask_time_length
												a__          =  mask_time_min_masks
												a__          =  mask_feature_prob
												a__          =  mask_feature_length
												a__          =  mask_feature_min_masks

												# ctc loss
												a__          =  ctc_loss_reduction
												a__          =  ctc_zero_infinity

												# sequence classification
												a__          =  use_weighted_layer_sum
												a__          =  classifier_proj_size





							@property
							def  lowercase		(   self:     Optional[Any]    ):

												'''simple docstring'''

												return functools.reduce(operator.mul				,		self.conv_stride				,		1    )










"""simple docstring"""







import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Optional[int]				,		__A:     Union[str, Any]				,		__A:     str				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Union[str, Any]    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=__A				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )
												a__          =  resnets

												if self.add_downsample:
																	a__          =  FlaxDownsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Dict				,		__A:     int				,		__A:     Dict				,		__A:     Optional[Any]=True    ):

												'''simple docstring'''

												a__          =  ()

												for resnet in self.resnets:
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	output_states += (hidden_states,)

												if self.add_downsample:
																	a__          =  self.downsamplers_a(__A    )
																	output_states += (hidden_states,)

												return hidden_states, output_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Optional[int]    ):

												'''simple docstring'''

												a__          =  []
												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.out_channels				,		n_heads=self.num_attention_heads				,		d_head=self.out_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		only_cross_attention=self.only_cross_attention				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

												a__          =  resnets
												a__          =  attentions

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     List[Any]				,		__A:     List[str]				,		__A:     Optional[Any]				,		__A:     Any=True    ):

												'''simple docstring'''

												for resnet, attn in zip(self.resnets				,		self.attentions    ):
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =True
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     str    ):

												'''simple docstring'''

												a__          =  []

												for i in range(self.num_layers    ):
																	a__          =  self.in_channels if (i == self.num_layers - 1) else self.out_channels
																	a__          =  self.prev_output_channel if i == 0 else self.out_channels

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=resnet_in_channels + res_skip_channels				,		out_channels=self.out_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets

												if self.add_upsample:
																	a__          =  FlaxUpsampleaD(self.out_channels				,		dtype=self.dtype    )





							def __call__(   self:     Tuple				,		__A:     Optional[Any]				,		__A:     Optional[Any]				,		__A:     Union[str, Any]				,		__A:     Dict=True    ):

												'''simple docstring'''

												for resnet in self.resnets:
																	# pop res hidden states
																	a__          =  res_hidden_states_tuple[-1]
																	a__          =  res_hidden_states_tuple[:-1]
																	a__          =  jnp.concatenate((hidden_states, res_hidden_states)				,		axis=-1    )

																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												if self.add_upsample:
																	a__          =  self.upsamplers_a(__A    )

												return hidden_states







class        _SCREAMING_SNAKE_CASE						(  nn.Module							):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE    =42
							_SCREAMING_SNAKE_CASE    =0.0
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =1
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =False
							_SCREAMING_SNAKE_CASE    =jnp.floataa




							def  lowercase		(   self:     Tuple    ):

												'''simple docstring'''

												a__          =  [
												    FlaxResnetBlockaD(
												        in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
												]

												a__          =  []

												for _ in range(self.num_layers    ):
																	a__          =  FlaxTransformeraDModel(
																	    in_channels=self.in_channels				,		n_heads=self.num_attention_heads				,		d_head=self.in_channels // self.num_attention_heads				,		depth=1				,		use_linear_projection=self.use_linear_projection				,		use_memory_efficient_attention=self.use_memory_efficient_attention				,		dtype=self.dtype				,		)
																	attentions.append(__A    )

																	a__          =  FlaxResnetBlockaD(
																	    in_channels=self.in_channels				,		out_channels=self.in_channels				,		dropout_prob=self.dropout				,		dtype=self.dtype				,		)
																	resnets.append(__A    )

												a__          =  resnets
												a__          =  attentions





							def __call__(   self:     Any				,		__A:     Optional[int]				,		__A:     int				,		__A:     Tuple				,		__A:     str=True    ):

												'''simple docstring'''

												a__          =  self.resnets[0](__A				,		__A    )
												for attn, resnet in zip(self.attentions				,		self.resnets[1:]    ):
																	a__          =  attn(__A				,		__A				,		deterministic=__A    )
																	a__          =  resnet(__A				,		__A				,		deterministic=__A    )

												return hidden_states


