import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map
def       __magic_name__(  __UpperCAmelCase		)						->    List[Tuple[int, ...]]:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		):
            for v in tree.values():
                  shapes.extend(_fetch_dims(__UpperCAmelCase		)		)
      elif isinstance(__UpperCAmelCase ,						(list, tuple)		):
            for t in tree:
                  shapes.extend(_fetch_dims(__UpperCAmelCase		)		)
      elif isinstance(__UpperCAmelCase ,						torch.Tensor		):
            shapes.append(tree.shape		)
      else:
            raise ValueError('''Not supported'''		)

      return shapes
@torch.jit.ignore
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    Tuple[int, ...]:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      for d in reversed(__UpperCAmelCase		):
            idx.append(flat_idx % d		)
            _lowerCamelCase			    =				flat_idx // d

      return tuple(reversed(__UpperCAmelCase		)		)
@torch.jit.ignore
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase = None ,						__UpperCAmelCase = None ,						)						->    List[Tuple[slice, ...]]:



      '''simple docstring'''

      def reduce_edge_list(__UpperCAmelCase		) -> None:
            _lowerCamelCase			    =				True
            for i in range(len(__UpperCAmelCase		)		):
                  _lowerCamelCase			    =				-1 * (i + 1)
                  l[reversed_idx] &= tally
                  _lowerCamelCase			    =				l[reversed_idx]

      if start_edges is None:
            _lowerCamelCase			    =				[s == 0 for s in start]
            reduce_edge_list(__UpperCAmelCase		)
      if end_edges is None:
            _lowerCamelCase			    =				[e == (d - 1) for e, d in zip(__UpperCAmelCase ,						__UpperCAmelCase		)]
            reduce_edge_list(__UpperCAmelCase		)

      # Base cases. Either start/end are empty and we're done, or the final,
      # one-dimensional tensor can be simply sliced
      if len(__UpperCAmelCase		) == 0:
            return [()]
      elif len(__UpperCAmelCase		) == 1:
            return [(slice(start[0] ,						end[0] + 1		),)]

      _lowerCamelCase			    =				[]
      _lowerCamelCase			    =				[]

      # Dimensions common to start and end can be selected directly
      for s, e in zip(__UpperCAmelCase ,						__UpperCAmelCase		):
            if s == e:
                  path_list.append(slice(__UpperCAmelCase ,						s + 1		)		)
            else:
                  break

      _lowerCamelCase			    =				tuple(__UpperCAmelCase		)
      _lowerCamelCase			    =				len(__UpperCAmelCase		)

      # start == end, and we're done
      if divergence_idx == len(__UpperCAmelCase		):
            return [path]

      def upper() -> Tuple[Tuple[slice, ...], ...]:
            assert start_edges is not None
            assert end_edges is not None

            _lowerCamelCase			    =				start[divergence_idx]
            return tuple(
                path + (slice(__UpperCAmelCase ,						sdi + 1		),) + s
                for s in _get_minimal_slice_set(
                    start[divergence_idx + 1 :] ,						[d - 1 for d in dims[divergence_idx + 1 :]] ,						dims[divergence_idx + 1 :] ,						start_edges=start_edges[divergence_idx + 1 :] ,						end_edges=[True for _ in end_edges[divergence_idx + 1 :]] ,						)		)

      def lower() -> Tuple[Tuple[slice, ...], ...]:
            assert start_edges is not None
            assert end_edges is not None

            _lowerCamelCase			    =				end[divergence_idx]
            return tuple(
                path + (slice(__UpperCAmelCase ,						edi + 1		),) + s
                for s in _get_minimal_slice_set(
                    [0 for _ in start[divergence_idx + 1 :]] ,						end[divergence_idx + 1 :] ,						dims[divergence_idx + 1 :] ,						start_edges=[True for _ in start_edges[divergence_idx + 1 :]] ,						end_edges=end_edges[divergence_idx + 1 :] ,						)		)

      # If both start and end are at the edges of the subtree rooted at
      # divergence_idx, we can just select the whole subtree at once
      if start_edges[divergence_idx] and end_edges[divergence_idx]:
            slices.append(path + (slice(start[divergence_idx] ,						end[divergence_idx] + 1		),)		)
      # If just start is at the edge, we can grab almost all of the subtree,
      # treating only the ragged bottom edge as an edge case
      elif start_edges[divergence_idx]:
            slices.append(path + (slice(start[divergence_idx] ,						end[divergence_idx]		),)		)
            slices.extend(lower()		)
      # Analogous to the previous case, but the top is ragged this time
      elif end_edges[divergence_idx]:
            slices.extend(upper()		)
            slices.append(path + (slice(start[divergence_idx] + 1 ,						end[divergence_idx] + 1		),)		)
      # If both sides of the range are ragged, we need to handle both sides
      # separately. If there's contiguous meat in between them, we can index it
      # in one big chunk
      else:
            slices.extend(upper()		)
            _lowerCamelCase			    =				end[divergence_idx] - start[divergence_idx]
            if middle_ground > 1:
                  slices.append(path + (slice(start[divergence_idx] + 1 ,						end[divergence_idx]		),)		)
            slices.extend(lower()		)

      return slices
@torch.jit.ignore
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase		)						->    torch.Tensor:



      '''simple docstring'''

      _lowerCamelCase			    =				t.shape[:no_batch_dims]
      _lowerCamelCase			    =				list(_flat_idx_to_idx(__UpperCAmelCase ,						__UpperCAmelCase		)		)
      # _get_minimal_slice_set is inclusive
      _lowerCamelCase			    =				list(_flat_idx_to_idx(flat_end - 1 ,						__UpperCAmelCase		)		)

      # Get an ordered list of slices to perform
      _lowerCamelCase			    =				_get_minimal_slice_set(
          __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						)

      _lowerCamelCase			    =				[t[s] for s in slices]

      return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]		) for s in sliced_tensors]		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase = False ,						__UpperCAmelCase = None ,						__UpperCAmelCase = False ,						)						->    Any:



      '''simple docstring'''

      if not (len(__UpperCAmelCase		) > 0):
            raise ValueError('''Must provide at least one input'''		)

      _lowerCamelCase			    =				[shape[:no_batch_dims] for shape in _fetch_dims(__UpperCAmelCase		)]
      _lowerCamelCase			    =				tuple([max(__UpperCAmelCase		) for s in zip(*__UpperCAmelCase		)]		)

      def _prep_inputs(__UpperCAmelCase		) -> torch.Tensor:
            if not low_mem:
                  if not sum(t.shape[:no_batch_dims]		) == no_batch_dims:
                        _lowerCamelCase			    =				t.expand(orig_batch_dims + t.shape[no_batch_dims:]		)
                  _lowerCamelCase			    =				t.reshape(-1 ,						*t.shape[no_batch_dims:]		)
            else:
                  _lowerCamelCase			    =				t.expand(orig_batch_dims + t.shape[no_batch_dims:]		)
            return t

      _lowerCamelCase			    =				tensor_tree_map(_prep_inputs ,						__UpperCAmelCase		)
      _lowerCamelCase			    =				None
      if _out is not None:
            _lowerCamelCase			    =				tensor_tree_map(lambda __UpperCAmelCase		: t.view([-1] + list(t.shape[no_batch_dims:]		)		) ,						_out		)

      _lowerCamelCase			    =				1
      for d in orig_batch_dims:
            flat_batch_dim *= d

      _lowerCamelCase			    =				flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

      def _select_chunk(__UpperCAmelCase		) -> torch.Tensor:
            return t[i : i + chunk_size] if t.shape[0] != 1 else t

      _lowerCamelCase			    =				0
      _lowerCamelCase			    =				prepped_outputs
      for _ in range(__UpperCAmelCase		):
            # Chunk the input
            if not low_mem:
                  _lowerCamelCase			    =				_select_chunk
            else:
                  _lowerCamelCase			    =				partial(
                      _chunk_slice ,						flat_start=__UpperCAmelCase ,						flat_end=min(__UpperCAmelCase ,						i + chunk_size		) ,						no_batch_dims=len(__UpperCAmelCase		) ,						)

            _lowerCamelCase			    =				tensor_tree_map(__UpperCAmelCase ,						__UpperCAmelCase		)

            # Run the layer on the chunk
            _lowerCamelCase			    =				layer(**__UpperCAmelCase		)

            # Allocate space for the output
            if out is None:
                  _lowerCamelCase			    =				tensor_tree_map(lambda __UpperCAmelCase		: t.new_zeros((flat_batch_dim,) + t.shape[1:]		) ,						__UpperCAmelCase		)

            # Put the chunk in its pre-allocated space
            if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		):

                  def assign(__UpperCAmelCase ,						__UpperCAmelCase		) -> None:
                        for k, v in da.items():
                              if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		):
                                    assign(__UpperCAmelCase ,						da[k]		)
                              else:
                                    if _add_into_out:
                                          v[i : i + chunk_size] += da[k]
                                    else:
                                          _lowerCamelCase			    =				da[k]

                  assign(__UpperCAmelCase ,						__UpperCAmelCase		)
            elif isinstance(__UpperCAmelCase ,						__UpperCAmelCase		):
                  for xa, xa in zip(__UpperCAmelCase ,						__UpperCAmelCase		):
                        if _add_into_out:
                              xa[i : i + chunk_size] += xa
                        else:
                              _lowerCamelCase			    =				xa
            elif isinstance(__UpperCAmelCase ,						torch.Tensor		):
                  if _add_into_out:
                        out[i : i + chunk_size] += output_chunk
                  else:
                        _lowerCamelCase			    =				output_chunk
            else:
                  raise ValueError('''Not supported'''		)

            i += chunk_size

      _lowerCamelCase			    =				tensor_tree_map(lambda __UpperCAmelCase		: t.view(orig_batch_dims + t.shape[1:]		) ,						__UpperCAmelCase		)

      return out




class  UpperCamelCase				:

     '''simple docstring'''

     def __init__(					self				,				A_ = 5_12				,				) ->				str:






           """simple docstring"""
           _lowerCamelCase			    =				max_chunk_size
           _lowerCamelCase			    =				None
           _lowerCamelCase			    =				None



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				int:






           """simple docstring"""
           logging.info('''Tuning chunk size...'''						)

           if min_chunk_size >= self.max_chunk_size:
                 return min_chunk_size

           _lowerCamelCase			    =				[2**l for l in range(int(math.log(self.max_chunk_size				,				2						)						) + 1						)]
           _lowerCamelCase			    =				[c for c in candidates if c > min_chunk_size]
           _lowerCamelCase			    =				[min_chunk_size] + candidates
           candidates[-1] += 4

           def test_chunk_size(A_						) -> bool:
                 try:
                       with torch.no_grad():
                             fn(*A_				,				chunk_size=A_						)
                       return True
                 except RuntimeError:
                       return False

           _lowerCamelCase			    =				0
           _lowerCamelCase			    =				len(A_						) - 1
           while i > min_viable_chunk_size_index:
                 _lowerCamelCase			    =				test_chunk_size(candidates[i]						)
                 if not viable:
                       _lowerCamelCase			    =				(min_viable_chunk_size_index + i) // 2
                 else:
                       _lowerCamelCase			    =				i
                       _lowerCamelCase			    =				(i + len(A_						) - 1) // 2

           return candidates[min_viable_chunk_size_index]



     def 		UpperCamelCase_						(					self				,				A_				,				A_						) ->				bool:






           """simple docstring"""
           _lowerCamelCase			    =				True
           for aa, aa in zip(A_				,				A_						):
                 assert type(A_						) == type(A_						)
                 if isinstance(A_				,				(list, tuple)						):
                       consistent &= self._compare_arg_caches(A_				,				A_						)
                 elif isinstance(A_				,				A_						):
                       _lowerCamelCase			    =				[v for _, v in sorted(aa.items()				,				key=lambda A_						: x[0]						)]
                       _lowerCamelCase			    =				[v for _, v in sorted(aa.items()				,				key=lambda A_						: x[0]						)]
                       consistent &= self._compare_arg_caches(A_				,				A_						)
                 else:
                       consistent &= aa == aa

           return consistent





     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_				,				) ->				int:






           """simple docstring"""
           _lowerCamelCase			    =				True
           _lowerCamelCase			    =				tree_map(lambda A_						: a.shape if isinstance(A_				,				torch.Tensor						) else a				,				A_				,				A_						)
           if self.cached_arg_data is not None:
                 # If args have changed shape/value, we need to re-tune
                 assert len(self.cached_arg_data						) == len(A_						)
                 _lowerCamelCase			    =				self._compare_arg_caches(self.cached_arg_data				,				A_						)
           else:
                 # Otherwise, we can reuse the precomputed value
                 _lowerCamelCase			    =				False

           if not consistent:
                 _lowerCamelCase			    =				self._determine_favorable_chunk_size(
                     A_				,				A_				,				A_				,				)
                 _lowerCamelCase			    =				arg_data

           assert self.cached_chunk_size is not None

           return self.cached_chunk_size
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


snake_case__				= {
    'configuration_whisper': ['WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'WhisperConfig', 'WhisperOnnxConfig'],
    'feature_extraction_whisper': ['WhisperFeatureExtractor'],
    'processing_whisper': ['WhisperProcessor'],
    'tokenization_whisper': ['WhisperTokenizer'],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= ['WhisperTokenizerFast']

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= [
         'WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'WhisperForConditionalGeneration',
         'WhisperModel',
         'WhisperPreTrainedModel',
         'WhisperForAudioClassification',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= [
         'TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'TFWhisperForConditionalGeneration',
         'TFWhisperModel',
         'TFWhisperPreTrainedModel',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__				= [
         'FlaxWhisperForConditionalGeneration',
         'FlaxWhisperModel',
         'FlaxWhisperPreTrainedModel',
         'FlaxWhisperForAudioClassification',
     ]


if TYPE_CHECKING:
     from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
     from .feature_extraction_whisper import WhisperFeatureExtractor
     from .processing_whisper import WhisperProcessor
     from .tokenization_whisper import WhisperTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_whisper_fast import WhisperTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_whisper import (
              WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
              WhisperForAudioClassification,
              WhisperForConditionalGeneration,
              WhisperModel,
              WhisperPreTrainedModel,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_whisper import (
              TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFWhisperForConditionalGeneration,
              TFWhisperModel,
              TFWhisperPreTrainedModel,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_whisper import (
              FlaxWhisperForAudioClassification,
              FlaxWhisperForConditionalGeneration,
              FlaxWhisperModel,
              FlaxWhisperPreTrainedModel,
          )

else:
     import sys

     snake_case__				= _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)