





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       UpperCamelCase__(			UpperCamelCase__     :			Union[dict, list, tuple, torch.Tensor]    )->List[Tuple[int, ...]]:
       A__		   =		[]
       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       UpperCamelCase__(			UpperCamelCase__     :			int					,	UpperCamelCase__     :			Tuple[int, ...]    )->Tuple[int, ...]:
       A__		   =		[]
       for d in reversed(UpperCamelCase__    ):
              idx.append(flat_idx % d    )
              A__		   =		flat_idx // d

       return tuple(reversed(UpperCamelCase__    )    )


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

       # 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(UpperCamelCase__     :			List[bool]    ) -> None:
              A__		   =		True
              for i in range(len(UpperCamelCase__    )    ):
                     A__		   =		-1 * (i + 1)
                     l[reversed_idx] &= tally
                     A__		   =		l[reversed_idx]

       if start_edges is None:
              A__		   =		[s == 0 for s in start]
              reduce_edge_list(UpperCamelCase__    )
       if end_edges is None:
              A__		   =		[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    ),)]

       A__		   =		[]
       A__		   =		[]

       # 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

       A__		   =		tuple(UpperCamelCase__    )
       A__		   =		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

              A__		   =		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

              A__		   =		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()    )
              A__		   =		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       UpperCamelCase__(			UpperCamelCase__     :			torch.Tensor					,	UpperCamelCase__     :			int					,	UpperCamelCase__     :			int					,	UpperCamelCase__     :			int    )->torch.Tensor:
       A__		   =		t.shape[:no_batch_dims]
       A__		   =		list(_flat_idx_to_idx(UpperCamelCase__					,	UpperCamelCase__    )    )
       # _get_minimal_slice_set is inclusive
       A__		   =		list(_flat_idx_to_idx(flat_end - 1					,	UpperCamelCase__    )    )

       # Get an ordered list of slices to perform
       A__		   =		_get_minimal_slice_set(
           UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	)

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

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


def       UpperCamelCase__(			UpperCamelCase__     :			Callable					,	UpperCamelCase__     :			Dict[str, Any]					,	UpperCamelCase__     :			int					,	UpperCamelCase__     :			int					,	UpperCamelCase__     :			bool = False					,	UpperCamelCase__     :			Any = None					,	UpperCamelCase__     :			bool = False					,	)->Any:
       if not (len(UpperCamelCase__    ) > 0):
              raise ValueError('''Must provide at least one input'''    )

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

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

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

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

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

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

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

              A__		   =		tensor_tree_map(UpperCamelCase__					,	UpperCamelCase__    )

              # Run the layer on the chunk
              A__		   =		layer(**UpperCamelCase__    )

              # Allocate space for the output
              if out is None:
                     A__		   =		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__     :			dict					,	UpperCamelCase__     :			dict    ) -> 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:
                                                 A__		   =		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:
                                   A__		   =		xa
              elif isinstance(UpperCamelCase__					,	torch.Tensor    ):
                     if _add_into_out:
                            out[i : i + chunk_size] += output_chunk
                     else:
                            A__		   =		output_chunk
              else:
                     raise ValueError('''Not supported'''    )

              i += chunk_size

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

       return out




class        SCREAMING_SNAKE_CASE__					:



 def __init__(				self,__lowerCamelCase = 512,):
        A__		   =		max_chunk_size
        A__		   =		None
        A__		   =		None



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					):
        logging.info('''Tuning chunk size...'''					)

        if min_chunk_size >= self.max_chunk_size:
               return min_chunk_size

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

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

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

        return candidates[min_viable_chunk_size_index]



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase					):
        A__		   =		True
        for aa, aa in zip(__lowerCamelCase,__lowerCamelCase					):
               assert type(__lowerCamelCase					) == type(__lowerCamelCase					)
               if isinstance(__lowerCamelCase,(list, tuple)					):
                      consistent &= self._compare_arg_caches(__lowerCamelCase,__lowerCamelCase					)
               elif isinstance(__lowerCamelCase,__lowerCamelCase					):
                      A__		   =		[v for _, v in sorted(aa.items(),key=lambda __lowerCamelCase					: x[0]					)]
                      A__		   =		[v for _, v in sorted(aa.items(),key=lambda __lowerCamelCase					: x[0]					)]
                      consistent &= self._compare_arg_caches(__lowerCamelCase,__lowerCamelCase					)
               else:
                      consistent &= aa == aa

        return consistent




 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__		   =		True
        A__		   =		tree_map(lambda __lowerCamelCase					: a.shape if isinstance(__lowerCamelCase,torch.Tensor					) else a,__lowerCamelCase,__lowerCamelCase					)
        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(__lowerCamelCase					)
               A__		   =		self._compare_arg_caches(self.cached_arg_data,__lowerCamelCase					)
        else:
               # Otherwise, we can reuse the precomputed value
               A__		   =		False

        if not consistent:
               A__		   =		self._determine_favorable_chunk_size(
                   __lowerCamelCase,__lowerCamelCase,__lowerCamelCase,)
               A__		   =		arg_data

        assert self.cached_chunk_size is not None

        return self.cached_chunk_size












a__: str							   =						'0.21.0'

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
       from .utils import rich





