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       lowerCAmelCase_     (		lowercase:							Union[dict, list, tuple, torch.Tensor]		)    ->				List[Tuple[int, ...]]:







 '''simple docstring'''
 _UpperCamelCase:     Union[str, Any]							= []
 if isinstance(lowercase  ,					lowercase		):
  for v in tree.values():
   shapes.extend(_fetch_dims(lowercase		)		)
 elif isinstance(lowercase  ,					(list, tuple)		):
  for t in tree:
   shapes.extend(_fetch_dims(lowercase		)		)
 elif isinstance(lowercase  ,					torch.Tensor		):
  shapes.append(tree.shape		)
 else:
  raise ValueError('''Not supported'''		)

 return shapes




@torch.jit.ignore
def       lowerCAmelCase_     (		lowercase:							int  ,					lowercase:							Tuple[int, ...]		)    ->				Tuple[int, ...]:







 '''simple docstring'''
 _UpperCamelCase:     Any							= []
 for d in reversed(lowercase		):
  idx.append(flat_idx % d		)
  _UpperCamelCase:     int							= flat_idx // d

 return tuple(reversed(lowercase		)		)




@torch.jit.ignore
def       lowerCAmelCase_     (		lowercase:							Sequence[int]  ,					lowercase:							Sequence[int]  ,					lowercase:							Sequence[int]  ,					lowercase:							Optional[Sequence[bool]] = None  ,					lowercase:							Optional[Sequence[bool]] = None  ,					)    ->				List[Tuple[slice, ...]]:







 '''simple docstring'''

 # start_edges and end_edges both indicate whether, starting from any given
 # dimension, the start/end index is at the top/bottom edge of the
 # corresponding tensor, modeled as a tree
 def reduce_edge_list(lowercase:							List[bool]		) -> None:
  _UpperCamelCase:     Optional[Any]							= True
  for i in range(len(lowercase		)		):
   _UpperCamelCase:     Optional[Any]							= -1 * (i + 1)
   l[reversed_idx] &= tally
   _UpperCamelCase:     Any							= l[reversed_idx]

 if start_edges is None:
  _UpperCamelCase:     Union[str, Any]							= [s == 0 for s in start]
  reduce_edge_list(lowercase		)
 if end_edges is None:
  _UpperCamelCase:     Any							= [e == (d - 1) for e, d in zip(lowercase  ,					lowercase		)]
  reduce_edge_list(lowercase		)

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

 _UpperCamelCase:     List[Tuple[slice, ...]]							= []
 _UpperCamelCase:     List[slice]							= []

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

 _UpperCamelCase:     Tuple[slice, ...]							= tuple(lowercase		)
 _UpperCamelCase:     Optional[Any]							= len(lowercase		)

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

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

  _UpperCamelCase:     List[Any]							= start[divergence_idx]
  return tuple(
      path + (slice(lowercase  ,					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

  _UpperCamelCase:     Tuple							= end[divergence_idx]
  return tuple(
      path + (slice(lowercase  ,					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()		)
  _UpperCamelCase:     Optional[int]							= 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       lowerCAmelCase_     (		lowercase:							torch.Tensor  ,					lowercase:							int  ,					lowercase:							int  ,					lowercase:							int		)    ->				torch.Tensor:







 '''simple docstring'''
 _UpperCamelCase:     int							= t.shape[:no_batch_dims]
 _UpperCamelCase:     Union[str, Any]							= list(_flat_idx_to_idx(lowercase  ,					lowercase		)		)
 # _get_minimal_slice_set is inclusive
 _UpperCamelCase:     List[str]							= list(_flat_idx_to_idx(flat_end - 1  ,					lowercase		)		)

 # Get an ordered list of slices to perform
 _UpperCamelCase:     Union[str, Any]							= _get_minimal_slice_set(
     lowercase  ,					lowercase  ,					lowercase  ,					)

 _UpperCamelCase:     str							= [t[s] for s in slices]

 return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]		) for s in sliced_tensors]		)




def       lowerCAmelCase_     (		lowercase:							Callable  ,					lowercase:							Dict[str, Any]  ,					lowercase:							int  ,					lowercase:							int  ,					lowercase:							bool = False  ,					lowercase:							Any = None  ,					lowercase:							bool = False  ,					)    ->				Any:







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

 _UpperCamelCase:     Dict							= [shape[:no_batch_dims] for shape in _fetch_dims(lowercase		)]
 _UpperCamelCase:     Optional[int]							= tuple([max(lowercase		) for s in zip(*lowercase		)]		)

 def _prep_inputs(lowercase:							torch.Tensor		) -> torch.Tensor:
  if not low_mem:
   if not sum(t.shape[:no_batch_dims]		) == no_batch_dims:
    _UpperCamelCase:     Optional[Any]							= t.expand(orig_batch_dims + t.shape[no_batch_dims:]		)
   _UpperCamelCase:     Tuple							= t.reshape(-1  ,					*t.shape[no_batch_dims:]		)
  else:
   _UpperCamelCase:     Dict							= t.expand(orig_batch_dims + t.shape[no_batch_dims:]		)
  return t

 _UpperCamelCase:     Dict[str, Any]							= tensor_tree_map(_prep_inputs  ,					lowercase		)
 _UpperCamelCase:     List[str]							= None
 if _out is not None:
  _UpperCamelCase:     Union[str, Any]							= tensor_tree_map(lambda lowercase		: t.view([-1] + list(t.shape[no_batch_dims:]		)		)  ,					_out		)

 _UpperCamelCase:     int							= 1
 for d in orig_batch_dims:
  flat_batch_dim *= d

 _UpperCamelCase:     int							= flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

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

 _UpperCamelCase:     Tuple							= 0
 _UpperCamelCase:     List[Any]							= prepped_outputs
 for _ in range(lowercase		):
  # Chunk the input
  if not low_mem:
   _UpperCamelCase:     Union[str, Any]							= _select_chunk
  else:
   _UpperCamelCase:     Dict							= partial(
       _chunk_slice  ,					flat_start=lowercase  ,					flat_end=min(lowercase  ,					i + chunk_size		)  ,					no_batch_dims=len(lowercase		)  ,					)

  _UpperCamelCase:     Dict[str, Any]							= tensor_tree_map(lowercase  ,					lowercase		)

  # Run the layer on the chunk
  _UpperCamelCase:     Union[str, Any]							= layer(**lowercase		)

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

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

   def assign(lowercase:							dict  ,					lowercase:							dict		) -> None:
    for k, v in da.items():
     if isinstance(lowercase  ,					lowercase		):
      assign(lowercase  ,					da[k]		)
     else:
      if _add_into_out:
       v[i : i + chunk_size] += da[k]
      else:
       _UpperCamelCase:     Union[str, Any]							= da[k]

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

  i += chunk_size

 _UpperCamelCase:     List[Any]							= tensor_tree_map(lambda lowercase		: t.view(orig_batch_dims + t.shape[1:]		)  ,					lowercase		)

 return out





class  __magic_name__  :





   """simple docstring"""




   def __init__(	self					:  int		,  _lowercase					:  int = 512		,  ):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= max_chunk_size
    _UpperCamelCase:     Optional[int]							= None
    _UpperCamelCase:     Optional[tuple]							= None



   def   lowerCAmelCase						(	self					:  int		,  _lowercase					:  Callable		,  _lowercase					:  tuple		,  _lowercase					:  int							):



    """simple docstring"""



    logging.info('''Tuning chunk size...'''							)

    if min_chunk_size >= self.max_chunk_size:
     return min_chunk_size

    _UpperCamelCase:     List[int]							= [2**l for l in range(int(math.log(self.max_chunk_size		,  2							)							) + 1							)]
    _UpperCamelCase:     str							= [c for c in candidates if c > min_chunk_size]
    _UpperCamelCase:     Union[str, Any]							= [min_chunk_size] + candidates
    candidates[-1] += 4

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

    _UpperCamelCase:     Union[str, Any]							= 0
    _UpperCamelCase:     Dict							= len(_lowercase							) - 1
    while i > min_viable_chunk_size_index:
     _UpperCamelCase:     Dict							= test_chunk_size(candidates[i]							)
     if not viable:
      _UpperCamelCase:     Dict							= (min_viable_chunk_size_index + i) // 2
     else:
      _UpperCamelCase:     int							= i
      _UpperCamelCase:     List[Any]							= (i + len(_lowercase							) - 1) // 2

    return candidates[min_viable_chunk_size_index]



   def   lowerCAmelCase						(	self					:  Union[str, Any]		,  _lowercase					:  Iterable		,  _lowercase					:  Iterable							):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= True
    for aa, aa in zip(_lowercase		,  _lowercase							):
     assert type(_lowercase							) == type(_lowercase							)
     if isinstance(_lowercase		,  (list, tuple)							):
      consistent &= self._compare_arg_caches(_lowercase		,  _lowercase							)
     elif isinstance(_lowercase		,  _lowercase							):
      _UpperCamelCase:     Dict							= [v for _, v in sorted(aa.items()		,  key=lambda _lowercase							: x[0]							)]
      _UpperCamelCase:     Optional[Any]							= [v for _, v in sorted(aa.items()		,  key=lambda _lowercase							: x[0]							)]
      consistent &= self._compare_arg_caches(_lowercase		,  _lowercase							)
     else:
      consistent &= aa == aa

    return consistent



   def   lowerCAmelCase						(	self					:  Union[str, Any]		,  _lowercase					:  Callable		,  _lowercase					:  tuple		,  _lowercase					:  int		,  ):



    """simple docstring"""



    _UpperCamelCase:     Dict							= True
    _UpperCamelCase:     tuple							= tree_map(lambda _lowercase							: a.shape if isinstance(_lowercase		,  torch.Tensor							) else a		,  _lowercase		,  _lowercase							)
    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(_lowercase							)
     _UpperCamelCase:     int							= self._compare_arg_caches(self.cached_arg_data		,  _lowercase							)
    else:
     # Otherwise, we can reuse the precomputed value
     _UpperCamelCase:     List[str]							= False

    if not consistent:
     _UpperCamelCase:     Optional[Any]							= self._determine_favorable_chunk_size(
         _lowercase		,  _lowercase		,  _lowercase		,  )
     _UpperCamelCase:     str							= arg_data

    assert self.cached_chunk_size is not None

    return self.cached_chunk_size
# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


UpperCAmelCase_			    =							[
    '''EAGER''',
    '''AOT_EAGER''',
    '''INDUCTOR''',
    '''NVFUSER''',
    '''AOT_NVFUSER''',
    '''AOT_CUDAGRAPHS''',
    '''OFI''',
    '''FX2TRT''',
    '''ONNXRT''',
    '''IPEX''',
]




def       lowerCAmelCase_     (		lowercase:							str  ,					lowercase:							Optional[Any]=None  ,					lowercase:							List[str]=None  ,					lowercase:							List[str]=None		)    ->				Tuple:







 '''simple docstring'''
 _UpperCamelCase:     int							= True
 while ask_again:
  _UpperCamelCase:     Any							= input(lowercase		)
  try:
   if default is not None and len(lowercase		) == 0:
    return default
   return convert_value(lowercase		) if convert_value is not None else result
  except Exception:
   if error_message is not None:
    print(lowercase		)




def       lowerCAmelCase_     (		lowercase:							List[Any]  ,					lowercase:							Union[str, Any]=[]  ,					lowercase:							Union[str, Any]=None  ,					lowercase:							Union[str, Any]=0		)    ->				List[Any]:







 '''simple docstring'''
 _UpperCamelCase:     List[str]							= BulletMenu(lowercase  ,					lowercase		)
 _UpperCamelCase:     List[Any]							= menu.run(default_choice=lowercase		)
 return convert_value(lowercase		) if convert_value is not None else result




def       lowerCAmelCase_     (		lowercase:							Dict		)    ->				Optional[int]:







 '''simple docstring'''
 _UpperCamelCase:     int							= int(lowercase		)
 return ComputeEnvironment(['''LOCAL_MACHINE''', '''AMAZON_SAGEMAKER'''][value]		)




def       lowerCAmelCase_     (		lowercase:							Tuple		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     str							= int(lowercase		)
 return DistributedType(['''NO''', '''MULTI_CPU''', '''MULTI_XPU''', '''MULTI_GPU''', '''MULTI_NPU''', '''TPU'''][value]		)




def       lowerCAmelCase_     (		lowercase:							Union[str, Any]		)    ->				Optional[int]:







 '''simple docstring'''
 _UpperCamelCase:     Tuple							= int(lowercase		)
 return DynamoBackend(DYNAMO_BACKENDS[value]		).value




def       lowerCAmelCase_     (		lowercase:							int		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     int							= int(lowercase		)
 return PrecisionType(['''no''', '''fp16''', '''bf16''', '''fp8'''][value]		)




def       lowerCAmelCase_     (		lowercase:							List[str]		)    ->				List[str]:







 '''simple docstring'''
 _UpperCamelCase:     Dict							= int(lowercase		)
 return SageMakerDistributedType(['''NO''', '''DATA_PARALLEL''', '''MODEL_PARALLEL'''][value]		)




def       lowerCAmelCase_     (		lowercase:							Any		)    ->				List[str]:







 '''simple docstring'''
 return {"yes": True, "no": False}[value.lower()]





class  __magic_name__  (		argparse.RawDescriptionHelpFormatter    ):





   """simple docstring"""




   def   lowerCAmelCase						(	self					:  str		,  _lowercase					:  Optional[int]		,  _lowercase					:  List[Any]		,  _lowercase					:  Dict		,  _lowercase					:  Any							):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= super()._format_usage(_lowercase		,  _lowercase		,  _lowercase		,  _lowercase							)
    _UpperCamelCase:     int							= usage.replace('''<command> [<args>] '''		,  ''''''							)
    return usage