'''simple docstring'''




from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block


@dataclass
class       UpperCAmelCase__  (       UpperCamelCase__			):






   """simple docstring"""
   A							:       Any     =    42



class       UpperCAmelCase__  (       nn.Module			):






   """simple docstring"""
   def __init__(self       ,				_a=3       ,				_a=3       ,				_a=("DownEncoderBlock2D",)       ,				_a=(64,)       ,				_a=2       ,				_a=32       ,				_a="silu"       ,				_a=True       ,				)    ->     List[str]:
        super().__init__()
        lowercase_					:       Any     				=					layers_per_block

        lowercase_					:       Optional[int]     				=					torch.nn.Convad(
            __A       ,				block_out_channels[0]       ,				kernel_size=3       ,				stride=1       ,				padding=1       ,				)

        lowercase_					:       List[Any]     				=					None
        lowercase_					:       Any     				=					nn.ModuleList([]				)

        # down
        lowercase_					:       int     				=					block_out_channels[0]
        for i, down_block_type in enumerate(__A				):
             lowercase_					:       Dict     				=					output_channel
             lowercase_					:       Tuple     				=					block_out_channels[i]
             lowercase_					:       Union[str, Any]     				=					i == len(__A				) - 1

             lowercase_					:       Any     				=					get_down_block(
                 __A       ,				num_layers=self.layers_per_block       ,				in_channels=__A       ,				out_channels=__A       ,				add_downsample=not is_final_block       ,				resnet_eps=1e-6       ,				downsample_padding=0       ,				resnet_act_fn=__A       ,				resnet_groups=__A       ,				attention_head_dim=__A       ,				temb_channels=__A       ,				)
             self.down_blocks.append(__A				)

        # mid
        lowercase_					:       Any     				=					UNetMidBlockaD(
            in_channels=block_out_channels[-1]       ,				resnet_eps=1e-6       ,				resnet_act_fn=__A       ,				output_scale_factor=1       ,				resnet_time_scale_shift='default'       ,				attention_head_dim=block_out_channels[-1]       ,				resnet_groups=__A       ,				temb_channels=__A       ,				)

        # out
        lowercase_					:       List[str]     				=					nn.GroupNorm(num_channels=block_out_channels[-1]       ,				num_groups=__A       ,				eps=1e-6				)
        lowercase_					:       str     				=					nn.SiLU()

        lowercase_					:       int     				=					2 * out_channels if double_z else out_channels
        lowercase_					:       Optional[int]     				=					nn.Convad(block_out_channels[-1]       ,				__A       ,				3       ,				padding=1				)

        lowercase_					:       Optional[Any]     				=					False





   def 		_lowerCamelCase (self       ,				_a				)    ->     Optional[int]:
        lowercase_					:       str     				=					x
        lowercase_					:       Dict     				=					self.conv_in(__A				)

        if self.training and self.gradient_checkpointing:

             def create_custom_forward(_a				):
                  def custom_forward(*_a				):
                       return module(*__A				)

                  return custom_forward

             # down
             if is_torch_version('>='       ,				'1.11.0'				):
                  for down_block in self.down_blocks:
                       lowercase_					:       List[str]     				=					torch.utils.checkpoint.checkpoint(
                           create_custom_forward(__A				)       ,				__A       ,				use_reentrant=__A				)
                  # middle
                  lowercase_					:       Optional[Any]     				=					torch.utils.checkpoint.checkpoint(
                      create_custom_forward(self.mid_block				)       ,				__A       ,				use_reentrant=__A				)
             else:
                  for down_block in self.down_blocks:
                       lowercase_					:       List[Any]     				=					torch.utils.checkpoint.checkpoint(create_custom_forward(__A				)       ,				__A				)
                  # middle
                  lowercase_					:       Tuple     				=					torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block				)       ,				__A				)

        else:
             # down
             for down_block in self.down_blocks:
                  lowercase_					:       Optional[Any]     				=					down_block(__A				)

             # middle
             lowercase_					:       Union[str, Any]     				=					self.mid_block(__A				)

        # post-process
        lowercase_					:       Dict     				=					self.conv_norm_out(__A				)
        lowercase_					:       int     				=					self.conv_act(__A				)
        lowercase_					:       Any     				=					self.conv_out(__A				)

        return sample



class       UpperCAmelCase__  (       nn.Module			):






   """simple docstring"""
   def __init__(self       ,				_a=3       ,				_a=3       ,				_a=("UpDecoderBlock2D",)       ,				_a=(64,)       ,				_a=2       ,				_a=32       ,				_a="silu"       ,				_a="group"       ,				)    ->     Optional[Any]:
        super().__init__()
        lowercase_					:       Tuple     				=					layers_per_block

        lowercase_					:       Optional[Any]     				=					nn.Convad(
            __A       ,				block_out_channels[-1]       ,				kernel_size=3       ,				stride=1       ,				padding=1       ,				)

        lowercase_					:       Any     				=					None
        lowercase_					:       str     				=					nn.ModuleList([]				)

        lowercase_					:       str     				=					in_channels if norm_type == 'spatial' else None

        # mid
        lowercase_					:       int     				=					UNetMidBlockaD(
            in_channels=block_out_channels[-1]       ,				resnet_eps=1e-6       ,				resnet_act_fn=__A       ,				output_scale_factor=1       ,				resnet_time_scale_shift='default' if norm_type == 'group' else norm_type       ,				attention_head_dim=block_out_channels[-1]       ,				resnet_groups=__A       ,				temb_channels=__A       ,				)

        # up
        lowercase_					:       Tuple     				=					list(reversed(__A				)				)
        lowercase_					:       int     				=					reversed_block_out_channels[0]
        for i, up_block_type in enumerate(__A				):
             lowercase_					:       Optional[Any]     				=					output_channel
             lowercase_					:       Tuple     				=					reversed_block_out_channels[i]

             lowercase_					:       Any     				=					i == len(__A				) - 1

             lowercase_					:       str     				=					get_up_block(
                 __A       ,				num_layers=self.layers_per_block + 1       ,				in_channels=__A       ,				out_channels=__A       ,				prev_output_channel=__A       ,				add_upsample=not is_final_block       ,				resnet_eps=1e-6       ,				resnet_act_fn=__A       ,				resnet_groups=__A       ,				attention_head_dim=__A       ,				temb_channels=__A       ,				resnet_time_scale_shift=__A       ,				)
             self.up_blocks.append(__A				)
             lowercase_					:       Union[str, Any]     				=					output_channel

        # out
        if norm_type == "spatial":
             lowercase_					:       List[Any]     				=					SpatialNorm(block_out_channels[0]       ,				__A				)
        else:
             lowercase_					:       Tuple     				=					nn.GroupNorm(num_channels=block_out_channels[0]       ,				num_groups=__A       ,				eps=1e-6				)
        lowercase_					:       List[str]     				=					nn.SiLU()
        lowercase_					:       Optional[Any]     				=					nn.Convad(block_out_channels[0]       ,				__A       ,				3       ,				padding=1				)

        lowercase_					:       Any     				=					False





   def 		_lowerCamelCase (self       ,				_a       ,				_a=None				)    ->     Optional[int]:
        lowercase_					:       Optional[int]     				=					z
        lowercase_					:       Dict     				=					self.conv_in(__A				)

        lowercase_					:       int     				=					next(iter(self.up_blocks.parameters()				)				).dtype
        if self.training and self.gradient_checkpointing:

             def create_custom_forward(_a				):
                  def custom_forward(*_a				):
                       return module(*__A				)

                  return custom_forward

             if is_torch_version('>='       ,				'1.11.0'				):
                  # middle
                  lowercase_					:       int     				=					torch.utils.checkpoint.checkpoint(
                      create_custom_forward(self.mid_block				)       ,				__A       ,				__A       ,				use_reentrant=__A				)
                  lowercase_					:       str     				=					sample.to(__A				)

                  # up
                  for up_block in self.up_blocks:
                       lowercase_					:       Tuple     				=					torch.utils.checkpoint.checkpoint(
                           create_custom_forward(__A				)       ,				__A       ,				__A       ,				use_reentrant=__A				)
             else:
                  # middle
                  lowercase_					:       List[Any]     				=					torch.utils.checkpoint.checkpoint(
                      create_custom_forward(self.mid_block				)       ,				__A       ,				__A				)
                  lowercase_					:       str     				=					sample.to(__A				)

                  # up
                  for up_block in self.up_blocks:
                       lowercase_					:       Union[str, Any]     				=					torch.utils.checkpoint.checkpoint(create_custom_forward(__A				)       ,				__A       ,				__A				)
        else:
             # middle
             lowercase_					:       Dict     				=					self.mid_block(__A       ,				__A				)
             lowercase_					:       Optional[Any]     				=					sample.to(__A				)

             # up
             for up_block in self.up_blocks:
                  lowercase_					:       Tuple     				=					up_block(__A       ,				__A				)

        # post-process
        if latent_embeds is None:
             lowercase_					:       Optional[int]     				=					self.conv_norm_out(__A				)
        else:
             lowercase_					:       List[Any]     				=					self.conv_norm_out(__A       ,				__A				)
        lowercase_					:       Union[str, Any]     				=					self.conv_act(__A				)
        lowercase_					:       Optional[int]     				=					self.conv_out(__A				)

        return sample



class       UpperCAmelCase__  (       nn.Module			):






   """simple docstring"""
   def __init__(self       ,				_a       ,				_a       ,				_a       ,				_a=None       ,				_a="random"       ,				_a=False       ,				_a=True				)    ->     Any:
        super().__init__()
        lowercase_					:       Optional[Any]     				=					n_e
        lowercase_					:       str     				=					vq_embed_dim
        lowercase_					:       Tuple     				=					beta
        lowercase_					:       Dict     				=					legacy

        lowercase_					:       List[Any]     				=					nn.Embedding(self.n_e       ,				self.vq_embed_dim				)
        self.embedding.weight.data.uniform_(-1.0 / self.n_e       ,				1.0 / self.n_e				)

        lowercase_					:       Union[str, Any]     				=					remap
        if self.remap is not None:
             self.register_buffer('used'       ,				torch.tensor(np.load(self.remap				)				)				)
             lowercase_					:       Any     				=					self.used.shape[0]
             lowercase_					:       Optional[Any]     				=					unknown_index  # "random" or "extra" or integer
             if self.unknown_index == "extra":
                  lowercase_					:       List[Any]     				=					self.re_embed
                  lowercase_					:       Optional[int]     				=					self.re_embed + 1
             print(
                 f'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
                 f'''Using {self.unknown_index} for unknown indices.'''				)
        else:
             lowercase_					:       str     				=					n_e

        lowercase_					:       Tuple     				=					sane_index_shape
   def 		_lowerCamelCase (self       ,				_a				)    ->     Union[str, Any]:
        lowercase_					:       str     				=					inds.shape
        assert len(__A				) > 1
        lowercase_					:       int     				=					inds.reshape(ishape[0]       ,				-1				)
        lowercase_					:       List[str]     				=					self.used.to(__A				)
        lowercase_					:       List[str]     				=					(inds[:, :, None] == used[None, None, ...]).long()
        lowercase_					:       Any     				=					match.argmax(-1				)
        lowercase_					:       Dict     				=					match.sum(2				) < 1
        if self.unknown_index == "random":
             lowercase_					:       Optional[int]     				=					torch.randint(0       ,				self.re_embed       ,				size=new[unknown].shape				).to(device=new.device				)
        else:
             lowercase_					:       Optional[int]     				=					self.unknown_index
        return new.reshape(__A				)
   def 		_lowerCamelCase (self       ,				_a				)    ->     List[Any]:
        lowercase_					:       int     				=					inds.shape
        assert len(__A				) > 1
        lowercase_					:       Optional[Any]     				=					inds.reshape(ishape[0]       ,				-1				)
        lowercase_					:       Union[str, Any]     				=					self.used.to(__A				)
        if self.re_embed > self.used.shape[0]:  # extra token
             lowercase_					:       Any     				=					0  # simply set to zero
        lowercase_					:       List[Any]     				=					torch.gather(used[None, :][inds.shape[0] * [0], :]       ,				1       ,				__A				)
        return back.reshape(__A				)
   def 		_lowerCamelCase (self       ,				_a				)    ->     Dict:
        lowercase_					:       str     				=					z.permute(0       ,				2       ,				3       ,				1				).contiguous()
        lowercase_					:       List[Any]     				=					z.view(-1       ,				self.vq_embed_dim				)

        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
        lowercase_					:       Optional[int]     				=					torch.argmin(torch.cdist(__A       ,				self.embedding.weight				)       ,				dim=1				)

        lowercase_					:       str     				=					self.embedding(__A				).view(z.shape				)
        lowercase_					:       Optional[Any]     				=					None
        lowercase_					:       Optional[int]     				=					None

        # compute loss for embedding
        if not self.legacy:
             lowercase_					:       Optional[Any]     				=					self.beta * torch.mean((z_q.detach() - z) ** 2				) + torch.mean((z_q - z.detach()) ** 2				)
        else:
             lowercase_					:       List[Any]     				=					torch.mean((z_q.detach() - z) ** 2				) + self.beta * torch.mean((z_q - z.detach()) ** 2				)

        # preserve gradients
        lowercase_					:       Any     				=					z + (z_q - z).detach()

        # reshape back to match original input shape
        lowercase_					:       Any     				=					z_q.permute(0       ,				3       ,				1       ,				2				).contiguous()

        if self.remap is not None:
             lowercase_					:       int     				=					min_encoding_indices.reshape(z.shape[0]       ,				-1				)  # add batch axis
             lowercase_					:       Any     				=					self.remap_to_used(__A				)
             lowercase_					:       List[Any]     				=					min_encoding_indices.reshape(-1       ,				1				)  # flatten

        if self.sane_index_shape:
             lowercase_					:       Tuple     				=					min_encoding_indices.reshape(z_q.shape[0]       ,				z_q.shape[2]       ,				z_q.shape[3]				)

        return z_q, loss, (perplexity, min_encodings, min_encoding_indices)





   def 		_lowerCamelCase (self       ,				_a       ,				_a				)    ->     List[Any]:



        if self.remap is not None:
             lowercase_					:       Dict     				=					indices.reshape(shape[0]       ,				-1				)  # add batch axis
             lowercase_					:       str     				=					self.unmap_to_all(__A				)
             lowercase_					:       Optional[Any]     				=					indices.reshape(-1				)  # flatten again

        # get quantized latent vectors
        lowercase_					:       List[Any]     				=					self.embedding(__A				)

        if shape is not None:
             lowercase_					:       str     				=					z_q.view(__A				)
             # reshape back to match original input shape
             lowercase_					:       Union[str, Any]     				=					z_q.permute(0       ,				3       ,				1       ,				2				).contiguous()

        return z_q



class       UpperCAmelCase__  (       UpperCamelCase__			):






   """simple docstring"""
   def __init__(self       ,				_a       ,				_a=False				)    ->     Any:
        lowercase_					:       int     				=					parameters
        lowercase_				,lowercase_					:       Tuple     				=					torch.chunk(__A       ,				2       ,				dim=1				)
        lowercase_					:       Dict     				=					torch.clamp(self.logvar       ,				-30.0       ,				20.0				)
        lowercase_					:       Union[str, Any]     				=					deterministic
        lowercase_					:       Optional[int]     				=					torch.exp(0.5 * self.logvar				)
        lowercase_					:       List[Any]     				=					torch.exp(self.logvar				)
        if self.deterministic:
             lowercase_					:       int     				=					torch.zeros_like(
                 self.mean       ,				device=self.parameters.device       ,				dtype=self.parameters.dtype				)
   def 		_lowerCamelCase (self       ,				_a = None				)    ->     torch.FloatTensor:
        lowercase_					:       List[str]     				=					randn_tensor(
            self.mean.shape       ,				generator=__A       ,				device=self.parameters.device       ,				dtype=self.parameters.dtype				)
        lowercase_					:       Dict     				=					self.mean + self.std * sample
        return x
   def 		_lowerCamelCase (self       ,				_a=None				)    ->     int:



        if self.deterministic:
             return torch.Tensor([0.0]				)
        else:
             if other is None:
                  return 0.5 * torch.sum(torch.pow(self.mean       ,				2				) + self.var - 1.0 - self.logvar       ,				dim=[1, 2, 3]				)
             else:
                  return 0.5 * torch.sum(
                      torch.pow(self.mean - other.mean       ,				2				) / other.var
                      + self.var / other.var
                      - 1.0
                      - self.logvar
                      + other.logvar       ,				dim=[1, 2, 3]       ,				)
   def 		_lowerCamelCase (self       ,				_a       ,				_a=[1, 2, 3]				)    ->     int:



        if self.deterministic:
             return torch.Tensor([0.0]				)
        lowercase_					:       Tuple     				=					np.log(2.0 * np.pi				)
        return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean       ,				2				) / self.var       ,				dim=__A				)





   def 		_lowerCamelCase (self				)    ->     Optional[int]:
        return self.mean





'''simple docstring'''




import numpy as np

from transformers import BatchFeature
from transformers.testing_utils import require_tf, require_torch

from .test_feature_extraction_common import FeatureExtractionSavingTestMixin


class       UpperCAmelCase__  (       _snake_case			):






   """simple docstring"""
   # to overwrite at feature extractactor specific tests
   A							:       List[str]     =    None
   A							:       Tuple     =    None
   @property
   def 		_lowerCamelCase (self				)    ->     Any:
        return self.feat_extract_tester.prepare_feat_extract_dict()
   def 		_lowerCamelCase (self				)    ->     str:
        lowercase_					:       Dict     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        self.assertTrue(hasattr(_a       ,				'feature_size'				)				)
        self.assertTrue(hasattr(_a       ,				'sampling_rate'				)				)
        self.assertTrue(hasattr(_a       ,				'padding_value'				)				)
   def 		_lowerCamelCase (self				)    ->     Dict:
        lowercase_					:       List[str]     				=					self.feat_extract_tester.prepare_inputs_for_common()
        lowercase_					:       Union[str, Any]     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        lowercase_					:       Union[str, Any]     				=					feat_extract.model_input_names[0]

        lowercase_					:       Union[str, Any]     				=					BatchFeature({input_name: speech_inputs}				)

        self.assertTrue(all(len(_a				) == len(_a				) for x, y in zip(_a       ,				processed_features[input_name]				)				)				)

        lowercase_					:       Optional[int]     				=					self.feat_extract_tester.prepare_inputs_for_common(equal_length=_a				)
        lowercase_					:       Tuple     				=					BatchFeature({input_name: speech_inputs}       ,				tensor_type='np'				)

        lowercase_					:       Dict     				=					processed_features[input_name]

        if len(batch_features_input.shape				) < 3:
             lowercase_					:       Optional[Any]     				=					batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]				), self.feat_extract_tester.feature_size)				)
   @require_torch
   def 		_lowerCamelCase (self				)    ->     Tuple:
        lowercase_					:       List[Any]     				=					self.feat_extract_tester.prepare_inputs_for_common(equal_length=_a				)
        lowercase_					:       Any     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        lowercase_					:       int     				=					feat_extract.model_input_names[0]

        lowercase_					:       Union[str, Any]     				=					BatchFeature({input_name: speech_inputs}       ,				tensor_type='pt'				)

        lowercase_					:       Tuple     				=					processed_features[input_name]

        if len(batch_features_input.shape				) < 3:
             lowercase_					:       List[str]     				=					batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]				), self.feat_extract_tester.feature_size)				)
   @require_tf
   def 		_lowerCamelCase (self				)    ->     Union[str, Any]:
        lowercase_					:       List[Any]     				=					self.feat_extract_tester.prepare_inputs_for_common(equal_length=_a				)
        lowercase_					:       int     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        lowercase_					:       Dict     				=					feat_extract.model_input_names[0]

        lowercase_					:       List[Any]     				=					BatchFeature({input_name: speech_inputs}       ,				tensor_type='tf'				)

        lowercase_					:       int     				=					processed_features[input_name]

        if len(batch_features_input.shape				) < 3:
             lowercase_					:       Union[str, Any]     				=					batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]				), self.feat_extract_tester.feature_size)				)
   def 		_lowerCamelCase (self       ,				_a=False				)    ->     Union[str, Any]:
        def _inputs_have_equal_length(_a				):
             lowercase_					:       Optional[Any]     				=					len(input[0]				)
             for input_slice in input[1:]:
                  if len(_a				) != length:
                       return False
             return True

        def _inputs_are_equal(_a       ,				_a				):
             if len(_a				) != len(_a				):
                  return False

             for input_slice_a, input_slice_a in zip(_a       ,				_a				):
                  if not np.allclose(np.asarray(_a				)       ,				np.asarray(_a				)       ,				atol=1e-3				):
                       return False
             return True

        lowercase_					:       int     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        lowercase_					:       str     				=					self.feat_extract_tester.prepare_inputs_for_common(numpify=_a				)
        lowercase_					:       Union[str, Any]     				=					feat_extract.model_input_names[0]

        lowercase_					:       Union[str, Any]     				=					BatchFeature({input_name: speech_inputs}				)

        lowercase_					:       Optional[Any]     				=					self.feat_extract_tester.seq_length_diff
        lowercase_					:       Union[str, Any]     				=					self.feat_extract_tester.max_seq_length + pad_diff
        lowercase_					:       str     				=					self.feat_extract_tester.min_seq_length
        lowercase_					:       Optional[Any]     				=					self.feat_extract_tester.batch_size
        lowercase_					:       Any     				=					self.feat_extract_tester.feature_size

        # test padding for List[int] + numpy
        lowercase_					:       List[str]     				=					feat_extract.pad(_a       ,				padding=_a				)
        lowercase_					:       Any     				=					input_a[input_name]

        lowercase_					:       str     				=					feat_extract.pad(_a       ,				padding='longest'				)
        lowercase_					:       List[str]     				=					input_a[input_name]

        lowercase_					:       Dict     				=					feat_extract.pad(_a       ,				padding='max_length'       ,				max_length=len(speech_inputs[-1]				)				)
        lowercase_					:       Any     				=					input_a[input_name]

        lowercase_					:       Any     				=					feat_extract.pad(_a       ,				padding='longest'       ,				return_tensors='np'				)
        lowercase_					:       str     				=					input_a[input_name]

        # max_length parameter has to be provided when setting `padding="max_length"`
        with self.assertRaises(_a				):
             feat_extract.pad(_a       ,				padding='max_length'				)[input_name]

        lowercase_					:       Union[str, Any]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=_a       ,				return_tensors='np'				)
        lowercase_					:       Any     				=					input_a[input_name]

        self.assertFalse(_inputs_have_equal_length(_a				)				)
        self.assertTrue(_inputs_have_equal_length(_a				)				)
        self.assertTrue(_inputs_have_equal_length(_a				)				)
        self.assertTrue(_inputs_are_equal(_a       ,				_a				)				)
        self.assertTrue(len(input_a[0]				) == pad_min_length				)
        self.assertTrue(len(input_a[1]				) == pad_min_length + pad_diff				)
        self.assertTrue(input_a.shape[:2] == (batch_size, len(input_a[0]				))				)
        self.assertTrue(input_a.shape[:2] == (batch_size, pad_max_length)				)

        if feature_size > 1:
             self.assertTrue(input_a.shape[2] == input_a.shape[2] == feature_size				)

        # test padding for `pad_to_multiple_of` for List[int] + numpy
        lowercase_					:       Tuple     				=					feat_extract.pad(_a       ,				pad_to_multiple_of=10				)
        lowercase_					:       Dict     				=					input_a[input_name]

        lowercase_					:       int     				=					feat_extract.pad(_a       ,				padding='longest'       ,				pad_to_multiple_of=10				)
        lowercase_					:       Union[str, Any]     				=					input_a[input_name]

        lowercase_					:       Optional[int]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				pad_to_multiple_of=10       ,				max_length=_a				)
        lowercase_					:       Tuple     				=					input_a[input_name]

        lowercase_					:       Optional[int]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				pad_to_multiple_of=10       ,				max_length=_a       ,				return_tensors='np'       ,				)
        lowercase_					:       Union[str, Any]     				=					input_a[input_name]

        self.assertTrue(all(len(_a				) % 10 == 0 for x in input_a				)				)
        self.assertTrue(_inputs_are_equal(_a       ,				_a				)				)

        lowercase_					:       str     				=					pad_max_length if pad_max_length % 10 == 0 else (pad_max_length // 10 + 1) * 10
        self.assertTrue(all(len(_a				) == expected_mult_pad_length for x in input_a				)				)
        self.assertEqual(input_a.shape[:2]       ,				(batch_size, expected_mult_pad_length)				)

        if feature_size > 1:
             self.assertTrue(input_a.shape[2] == feature_size				)

        # Check padding value is correct
        lowercase_					:       Optional[int]     				=					(np.ones(self.feat_extract_tester.feature_size				) * feat_extract.padding_value).sum()
        self.assertTrue(
            abs(np.asarray(input_a[0]				)[pad_min_length:].sum() - padding_vector_sum * (pad_max_length - pad_min_length)				)
            < 1e-3				)
        self.assertTrue(
            abs(
                np.asarray(input_a[1]				)[pad_min_length + pad_diff :].sum()
                - padding_vector_sum * (pad_max_length - pad_min_length - pad_diff)				)
            < 1e-3				)
        self.assertTrue(
            abs(
                np.asarray(input_a[2]				)[pad_min_length + 2 * pad_diff :].sum()
                - padding_vector_sum * (pad_max_length - pad_min_length - 2 * pad_diff)				)
            < 1e-3				)
        self.assertTrue(
            abs(input_a[0, pad_min_length:].sum() - padding_vector_sum * (pad_max_length - pad_min_length)				) < 1e-3				)
        self.assertTrue(
            abs(input_a[0, pad_min_length:].sum() - padding_vector_sum * (expected_mult_pad_length - pad_min_length)				)
            < 1e-3				)
   def 		_lowerCamelCase (self       ,				_a=False				)    ->     Union[str, Any]:
        def _inputs_have_equal_length(_a				):
             lowercase_					:       str     				=					len(input[0]				)
             for input_slice in input[1:]:
                  if len(_a				) != length:
                       return False
             return True

        def _inputs_are_equal(_a       ,				_a				):
             if len(_a				) != len(_a				):
                  return False

             for input_slice_a, input_slice_a in zip(_a       ,				_a				):
                  if not np.allclose(np.asarray(_a				)       ,				np.asarray(_a				)       ,				atol=1e-3				):
                       return False
             return True

        lowercase_					:       Any     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        lowercase_					:       List[Any]     				=					self.feat_extract_tester.prepare_inputs_for_common(numpify=_a				)
        lowercase_					:       List[str]     				=					feat_extract.model_input_names[0]

        lowercase_					:       str     				=					BatchFeature({input_name: speech_inputs}				)

        # truncate to smallest
        lowercase_					:       Union[str, Any]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[0]				)       ,				truncation=_a				)
        lowercase_					:       List[str]     				=					input_a[input_name]

        lowercase_					:       Optional[int]     				=					feat_extract.pad(_a       ,				padding='max_length'       ,				max_length=len(speech_inputs[0]				)				)
        lowercase_					:       Union[str, Any]     				=					input_a[input_name]

        self.assertTrue(_inputs_have_equal_length(_a				)				)
        self.assertFalse(_inputs_have_equal_length(_a				)				)

        # truncate to smallest with np
        lowercase_					:       List[str]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[0]				)       ,				return_tensors='np'       ,				truncation=_a       ,				)
        lowercase_					:       List[str]     				=					input_a[input_name]

        lowercase_					:       List[Any]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[0]				)       ,				return_tensors='np'				)
        lowercase_					:       List[Any]     				=					input_a[input_name]

        self.assertTrue(_inputs_have_equal_length(_a				)				)
        self.assertTrue(input_a.shape[1] == len(speech_inputs[0]				)				)

        # since truncation forces padding to be smaller than longest input
        # function can't return `np.ndarray`, but has to return list
        self.assertFalse(_inputs_have_equal_length(_a				)				)

        # truncate to middle
        lowercase_					:       int     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[1]				)       ,				truncation=_a       ,				return_tensors='np'       ,				)
        lowercase_					:       Union[str, Any]     				=					input_a[input_name]

        lowercase_					:       str     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[1]				)       ,				truncation=_a				)
        lowercase_					:       int     				=					input_a[input_name]

        lowercase_					:       Dict     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[1]				)       ,				return_tensors='np'				)
        lowercase_					:       Optional[int]     				=					input_a[input_name]

        self.assertTrue(input_a.shape[1] == len(speech_inputs[1]				)				)
        self.assertTrue(_inputs_have_equal_length(_a				)				)
        self.assertTrue(_inputs_have_equal_length(_a				)				)
        self.assertTrue(_inputs_are_equal(_a       ,				_a				)				)

        # since truncation forces padding to be smaller than longest input
        # function can't return `np.ndarray`, but has to return list
        self.assertFalse(_inputs_have_equal_length(_a				)				)
        self.assertTrue(len(input_a[-1]				) == len(speech_inputs[-1]				)				)

        # padding has to be max_length when setting `truncation=True`
        with self.assertRaises(_a				):
             feat_extract.pad(_a       ,				truncation=_a				)[input_name]

        # padding has to be max_length when setting `truncation=True`
        with self.assertRaises(_a				):
             feat_extract.pad(_a       ,				padding='longest'       ,				truncation=_a				)[input_name]

        # padding has to be max_length when setting `truncation=True`
        with self.assertRaises(_a				):
             feat_extract.pad(_a       ,				padding='longest'       ,				truncation=_a				)[input_name]

        # max_length parameter has to be provided when setting `truncation=True` and padding="max_length"
        with self.assertRaises(_a				):
             feat_extract.pad(_a       ,				padding='max_length'       ,				truncation=_a				)[input_name]

        # test truncation for `pad_to_multiple_of` for List[int] + numpy
        lowercase_					:       str     				=					12
        lowercase_					:       List[Any]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[0]				)       ,				pad_to_multiple_of=_a       ,				truncation=_a       ,				)
        lowercase_					:       Any     				=					input_a[input_name]

        lowercase_					:       List[Any]     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=len(speech_inputs[0]				)       ,				pad_to_multiple_of=_a       ,				)
        lowercase_					:       str     				=					input_a[input_name]

        # retrieve expected_length as multiple of pad_to_multiple_of
        lowercase_					:       Dict     				=					len(speech_inputs[0]				)
        if expected_length % pad_to_multiple_of != 0:
             lowercase_					:       List[Any]     				=					((len(speech_inputs[0]				) // pad_to_multiple_of) + 1) * pad_to_multiple_of

        self.assertTrue(len(input_a[0]				) == expected_length				)
        self.assertTrue(_inputs_have_equal_length(_a				)				)
        self.assertFalse(_inputs_have_equal_length(_a				)				)
   def 		_lowerCamelCase (self				)    ->     str:
        self._check_padding(numpify=_a				)
   def 		_lowerCamelCase (self				)    ->     Optional[int]:
        self._check_padding(numpify=_a				)
   def 		_lowerCamelCase (self				)    ->     Any:
        self._check_truncation(numpify=_a				)
   def 		_lowerCamelCase (self				)    ->     Any:
        self._check_truncation(numpify=_a				)
   @require_torch
   def 		_lowerCamelCase (self				)    ->     int:
        lowercase_					:       Tuple     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        lowercase_					:       List[Any]     				=					self.feat_extract_tester.prepare_inputs_for_common()
        lowercase_					:       Optional[Any]     				=					feat_extract.model_input_names[0]

        lowercase_					:       Dict     				=					BatchFeature({input_name: speech_inputs}				)

        lowercase_					:       Tuple     				=					feat_extract.pad(_a       ,				padding='longest'       ,				return_tensors='np'				)[input_name]
        lowercase_					:       List[Any]     				=					feat_extract.pad(_a       ,				padding='longest'       ,				return_tensors='pt'				)[input_name]

        self.assertTrue(abs(input_np.astype(np.floataa				).sum() - input_pt.numpy().astype(np.floataa				).sum()				) < 1e-2				)
   @require_tf
   def 		_lowerCamelCase (self				)    ->     List[Any]:
        lowercase_					:       Any     				=					self.feature_extraction_class(**self.feat_extract_dict				)
        lowercase_					:       Optional[int]     				=					self.feat_extract_tester.prepare_inputs_for_common()
        lowercase_					:       List[Any]     				=					feat_extract.model_input_names[0]

        lowercase_					:       Tuple     				=					BatchFeature({input_name: speech_inputs}				)

        lowercase_					:       List[Any]     				=					feat_extract.pad(_a       ,				padding='longest'       ,				return_tensors='np'				)[input_name]
        lowercase_					:       Optional[int]     				=					feat_extract.pad(_a       ,				padding='longest'       ,				return_tensors='tf'				)[input_name]

        self.assertTrue(abs(input_np.astype(np.floataa				).sum() - input_tf.numpy().astype(np.floataa				).sum()				) < 1e-2				)
   def 		_lowerCamelCase (self				)    ->     List[str]:
        lowercase_					:       Any     				=					self.feat_extract_dict
        lowercase_					:       List[Any]     				=					True
        lowercase_					:       Tuple     				=					self.feature_extraction_class(**_a				)
        lowercase_					:       Optional[Any]     				=					self.feat_extract_tester.prepare_inputs_for_common()
        lowercase_					:       str     				=					[len(_a				) for x in speech_inputs]
        lowercase_					:       Union[str, Any]     				=					feat_extract.model_input_names[0]

        lowercase_					:       List[Any]     				=					BatchFeature({input_name: speech_inputs}				)

        lowercase_					:       str     				=					feat_extract.pad(_a       ,				padding='longest'       ,				return_tensors='np'				)
        self.assertIn('attention_mask'       ,				_a				)
        self.assertListEqual(list(processed.attention_mask.shape				)       ,				list(processed[input_name].shape[:2]				)				)
        self.assertListEqual(processed.attention_mask.sum(-1				).tolist()       ,				_a				)





   def 		_lowerCamelCase (self				)    ->     List[Any]:
        lowercase_					:       Union[str, Any]     				=					self.feat_extract_dict
        lowercase_					:       List[Any]     				=					True
        lowercase_					:       List[Any]     				=					self.feature_extraction_class(**_a				)
        lowercase_					:       List[Any]     				=					self.feat_extract_tester.prepare_inputs_for_common()
        lowercase_					:       Optional[int]     				=					[len(_a				) for x in speech_inputs]
        lowercase_					:       int     				=					feat_extract.model_input_names[0]

        lowercase_					:       Tuple     				=					BatchFeature({input_name: speech_inputs}				)
        lowercase_					:       Optional[Any]     				=					min(_a				)

        lowercase_					:       int     				=					feat_extract.pad(
            _a       ,				padding='max_length'       ,				max_length=_a       ,				truncation=_a       ,				return_tensors='np'				)
        self.assertIn('attention_mask'       ,				_a				)
        self.assertListEqual(
            list(processed_pad.attention_mask.shape				)       ,				[processed_pad[input_name].shape[0], max_length]				)
        self.assertListEqual(
            processed_pad.attention_mask[:, :max_length].sum(-1				).tolist()       ,				[max_length for x in speech_inputs]				)




