'''simple docstring'''







import baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler




class _SCREAMING_SNAKE_CASE:




						def __init__(    self :       str							,							UpperCamelCase_ :       Dict      )     ->				str:
													if isinstance(UpperCamelCase_							,							UpperCamelCase_      ):
																				# Don't modify user's data should they want to reuse it (e.g. in tests), because once we
																				# modified it, it will not be accepted here again, since `auto` values would have been overridden
																				SCREAMING_SNAKE_CASE__     :str						    = deepcopy(UpperCamelCase_      )
													elif os.path.exists(UpperCamelCase_      ):
																				with io.open(UpperCamelCase_							,							'r'							,							encoding='utf-8'      ) as f:
																											SCREAMING_SNAKE_CASE__     :Dict						    = json.load(UpperCamelCase_      )
													else:
																				try:
																											SCREAMING_SNAKE_CASE__     :Tuple						    = baseaa.urlsafe_baadecode(UpperCamelCase_      ).decode('utf-8'      )
																											SCREAMING_SNAKE_CASE__     :Optional[Any]						    = json.loads(UpperCamelCase_      )
																				except (UnicodeDecodeError, AttributeError, ValueError):
																											raise ValueError(
																											    f'''Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}'''      )

													SCREAMING_SNAKE_CASE__     :str						    = config

													self.set_stage_and_offload()




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Optional[int]:
													# zero stage - this is done as early as possible, before model is created, to allow
													# ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
													# during ``zero.Init()`` which needs to know the dtype, and some other hparams.
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.get_value('zero_optimization.stage'							,							-1      )

													# offload
													SCREAMING_SNAKE_CASE__     :Tuple						    = False
													if self.is_zeroa() or self.is_zeroa():
																				SCREAMING_SNAKE_CASE__     :Dict						    = set(['cpu', 'nvme']      )
																				SCREAMING_SNAKE_CASE__     :Optional[int]						    = set(
																				    [
																				        self.get_value('zero_optimization.offload_optimizer.device'      ),
																				        self.get_value('zero_optimization.offload_param.device'      ),
																				    ]      )
																				if len(offload_devices & offload_devices_valid      ) > 0:
																											SCREAMING_SNAKE_CASE__     :Optional[int]						    = True




						def      __lowerCamelCase       (    self :       Optional[int]							,							UpperCamelCase_ :       List[Any]      )     ->				Optional[int]:
													SCREAMING_SNAKE_CASE__     :Tuple						    = self.config

													# find the config node of interest if it exists
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = ds_key_long.split('.'      )
													SCREAMING_SNAKE_CASE__     :Tuple						    = nodes.pop()
													for node in nodes:
																				SCREAMING_SNAKE_CASE__     :int						    = config.get(UpperCamelCase_      )
																				if config is None:
																											return None, ds_key

													return config, ds_key




						def      __lowerCamelCase       (    self :       Union[str, Any]							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       Union[str, Any]=None      )     ->				Any:
													SCREAMING_SNAKE_CASE__    ,      SCREAMING_SNAKE_CASE__     :Optional[Any]						    = self.find_config_node(UpperCamelCase_      )
													if config is None:
																				return default
													return config.get(UpperCamelCase_							,							UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       str							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Optional[Any]=False      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Tuple						    = self.config

													# find the config node of interest if it exists
													SCREAMING_SNAKE_CASE__     :Any						    = ds_key_long.split('.'      )
													for node in nodes:
																				SCREAMING_SNAKE_CASE__     :Any						    = config
																				SCREAMING_SNAKE_CASE__     :List[str]						    = config.get(UpperCamelCase_      )
																				if config is None:
																											if must_exist:
																																		raise ValueError(f'''Can\'t find {ds_key_long} entry in the config: {self.config}'''      )
																											else:
																																		return

        # if found remove it
													if parent_config is not None:
																				parent_config.pop(UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       List[Any]							,							UpperCamelCase_ :       Union[str, Any]      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = self.get_value(UpperCamelCase_      )
													return False if value is None else bool(UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Dict							,							UpperCamelCase_ :       List[str]      )     ->				Tuple:
													SCREAMING_SNAKE_CASE__     :List[str]						    = self.get_value(UpperCamelCase_      )
													return False if value is None else not bool(UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       int      )     ->				List[Any]:
													return self._stage == 2




						def      __lowerCamelCase       (    self :       int      )     ->				Union[str, Any]:
													return self._stage == 3




						def      __lowerCamelCase       (    self :       str      )     ->				List[str]:
													return self._offload




class _SCREAMING_SNAKE_CASE:




						def __init__(    self :       Tuple							,							UpperCamelCase_ :       List[Any]      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :List[str]						    = engine




						def      __lowerCamelCase       (    self :       Optional[Any]							,							UpperCamelCase_ :       Tuple							,							**UpperCamelCase_ :       Optional[Any]      )     ->				int:
													# runs backpropagation and handles mixed precision
													self.engine.backward(UpperCamelCase_							,							**UpperCamelCase_      )

													# Deepspeed's `engine.step` performs the following operations:
													# - gradient accumulation check
													# - gradient clipping
													# - optimizer step
													# - zero grad
													# - checking overflow
													# - lr_scheduler step (only if engine.lr_scheduler is not None)
													self.engine.step()
													# and this plugin overrides the above calls with no-ops when Accelerate runs under
													# Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
													# training loop that works transparently under many training regimes.




class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE					):




						def __init__(    self :       List[str]							,							UpperCamelCase_ :       List[str]      )     ->				Any:
													super().__init__(UpperCamelCase_							,							device_placement=UpperCamelCase_							,							scaler=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :Tuple						    = hasattr(self.optimizer							,							'overflow'      )




						def      __lowerCamelCase       (    self :       List[str]							,							UpperCamelCase_ :       str=None      )     ->				List[str]:
													pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				Tuple:
													pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




						@property
						def      __lowerCamelCase       (    self :       Optional[Any]      )     ->				str:
													if self.__has_overflow__:
																				return self.optimizer.overflow
													return False




class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE					):




						def __init__(    self :       Optional[Any]							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Dict      )     ->				Dict:
													super().__init__(UpperCamelCase_							,							UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Optional[int]:
													pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed




class _SCREAMING_SNAKE_CASE:




						def __init__(    self :       Dict							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       int=0.001							,							UpperCamelCase_ :       Optional[Any]=0							,							**UpperCamelCase_ :       List[Any]      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = params
													SCREAMING_SNAKE_CASE__     :Dict						    = lr
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = weight_decay
													SCREAMING_SNAKE_CASE__     :str						    = kwargs




class _SCREAMING_SNAKE_CASE:




						def __init__(    self :       Optional[int]							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       Any=None							,							UpperCamelCase_ :       Dict=0							,							**UpperCamelCase_ :       Any      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :str						    = optimizer
													SCREAMING_SNAKE_CASE__     :List[Any]						    = total_num_steps
													SCREAMING_SNAKE_CASE__     :List[str]						    = warmup_num_steps
													SCREAMING_SNAKE_CASE__     :Dict						    = kwargs







'''simple docstring'''







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


							'''simple docstring'''






							SCREAMING_SNAKE_CASE__     :Optional[int]						    = []
							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    lowerCamelCase       (			UpperCAmelCase__				:  int				,						UpperCAmelCase__				:  Tuple[int, ...]		)			->	Tuple[int, ...]:


							'''simple docstring'''






							SCREAMING_SNAKE_CASE__     :Dict						    = []
							for d in reversed(UpperCAmelCase__		):
														idx.append(flat_idx % d		)
														SCREAMING_SNAKE_CASE__     :Optional[int]						    = flat_idx // d

							return tuple(reversed(UpperCAmelCase__		)		)


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


							'''simple docstring'''






							def reduce_edge_list(UpperCAmelCase__				:  List[bool]		) -> None:
														SCREAMING_SNAKE_CASE__     :Optional[int]						    = True
														for i in range(len(UpperCAmelCase__		)		):
																					SCREAMING_SNAKE_CASE__     :List[Any]						    = -1 * (i + 1)
																					l[reversed_idx] &= tally
																					SCREAMING_SNAKE_CASE__     :Tuple						    = l[reversed_idx]

							if start_edges is None:
														SCREAMING_SNAKE_CASE__     :Optional[int]						    = [s == 0 for s in start]
														reduce_edge_list(UpperCAmelCase__		)
							if end_edges is None:
														SCREAMING_SNAKE_CASE__     :Dict						    = [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		),)]

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

							# 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

							SCREAMING_SNAKE_CASE__     :Tuple[slice, ...]						    = tuple(UpperCAmelCase__		)
							SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = 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

														SCREAMING_SNAKE_CASE__     :Optional[Any]						    = 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

														SCREAMING_SNAKE_CASE__     :Any						    = 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()		)
														SCREAMING_SNAKE_CASE__     :str						    = 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       (			UpperCAmelCase__				:  torch.Tensor				,						UpperCAmelCase__				:  int				,						UpperCAmelCase__				:  int				,						UpperCAmelCase__				:  int		)			->	torch.Tensor:


							'''simple docstring'''






							SCREAMING_SNAKE_CASE__     :Optional[Any]						    = t.shape[:no_batch_dims]
							SCREAMING_SNAKE_CASE__     :int						    = list(_flat_idx_to_idx(UpperCAmelCase__				,						UpperCAmelCase__		)		)
							# _get_minimal_slice_set is inclusive
							SCREAMING_SNAKE_CASE__     :Optional[Any]						    = list(_flat_idx_to_idx(flat_end - 1				,						UpperCAmelCase__		)		)

							# Get an ordered list of slices to perform
							SCREAMING_SNAKE_CASE__     :List[Any]						    = _get_minimal_slice_set(
							    UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				,						)

							SCREAMING_SNAKE_CASE__     :Dict						    = [t[s] for s in slices]

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


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


							'''simple docstring'''






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

							SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = [shape[:no_batch_dims] for shape in _fetch_dims(UpperCAmelCase__		)]
							SCREAMING_SNAKE_CASE__     :Optional[int]						    = 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:
																												SCREAMING_SNAKE_CASE__     :Tuple						    = t.expand(orig_batch_dims + t.shape[no_batch_dims:]		)
																					SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = t.reshape(-1				,						*t.shape[no_batch_dims:]		)
														else:
																					SCREAMING_SNAKE_CASE__     :Any						    = t.expand(orig_batch_dims + t.shape[no_batch_dims:]		)
														return t

							SCREAMING_SNAKE_CASE__     :Dict[str, Any]						    = tensor_tree_map(_prep_inputs				,						UpperCAmelCase__		)
							SCREAMING_SNAKE_CASE__     :Dict						    = None
							if _out is not None:
														SCREAMING_SNAKE_CASE__     :Optional[Any]						    = tensor_tree_map(lambda UpperCAmelCase__		: t.view([-1] + list(t.shape[no_batch_dims:]		)		)				,						_out		)

							SCREAMING_SNAKE_CASE__     :Optional[int]						    = 1
							for d in orig_batch_dims:
														flat_batch_dim *= d

							SCREAMING_SNAKE_CASE__     :str						    = 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

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

														SCREAMING_SNAKE_CASE__     :Dict[str, Any]						    = tensor_tree_map(UpperCAmelCase__				,						UpperCAmelCase__		)

														# Run the layer on the chunk
														SCREAMING_SNAKE_CASE__     :Any						    = layer(**UpperCAmelCase__		)

														# Allocate space for the output
														if out is None:
																					SCREAMING_SNAKE_CASE__     :Optional[Any]						    = 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:
																																																	SCREAMING_SNAKE_CASE__     :str						    = 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:
																																			SCREAMING_SNAKE_CASE__     :int						    = xa
														elif isinstance(UpperCAmelCase__				,						torch.Tensor		):
																					if _add_into_out:
																												out[i : i + chunk_size] += output_chunk
																					else:
																												SCREAMING_SNAKE_CASE__     :List[Any]						    = output_chunk
														else:
																					raise ValueError('Not supported'		)

														i += chunk_size

							SCREAMING_SNAKE_CASE__     :str						    = tensor_tree_map(lambda UpperCAmelCase__		: t.view(orig_batch_dims + t.shape[1:]		)				,						UpperCAmelCase__		)

							return out





class _SCREAMING_SNAKE_CASE:




						def __init__(    self :       Any							,							UpperCamelCase_ :       int = 5_12							,							)     ->				int:
													SCREAMING_SNAKE_CASE__     :int						    = max_chunk_size
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = None
													SCREAMING_SNAKE_CASE__     :Optional[tuple]						    = None




						def      __lowerCamelCase       (    self :       List[Any]							,							UpperCamelCase_ :       Callable							,							UpperCamelCase_ :       tuple							,							UpperCamelCase_ :       int      )     ->				int:
													logging.info('Tuning chunk size...'      )

													if min_chunk_size >= self.max_chunk_size:
																				return min_chunk_size

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

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

													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = 0
													SCREAMING_SNAKE_CASE__     :List[str]						    = len(UpperCamelCase_      ) - 1
													while i > min_viable_chunk_size_index:
																				SCREAMING_SNAKE_CASE__     :int						    = test_chunk_size(candidates[i]      )
																				if not viable:
																											SCREAMING_SNAKE_CASE__     :List[Any]						    = (min_viable_chunk_size_index + i) // 2
																				else:
																											SCREAMING_SNAKE_CASE__     :int						    = i
																											SCREAMING_SNAKE_CASE__     :List[str]						    = (i + len(UpperCamelCase_      ) - 1) // 2

													return candidates[min_viable_chunk_size_index]




						def      __lowerCamelCase       (    self :       List[str]							,							UpperCamelCase_ :       Iterable							,							UpperCamelCase_ :       Iterable      )     ->				bool:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = True
													for aa, aa in zip(UpperCamelCase_							,							UpperCamelCase_      ):
																				assert type(UpperCamelCase_      ) == type(UpperCamelCase_      )
																				if isinstance(UpperCamelCase_							,							(list, tuple)      ):
																											consistent &= self._compare_arg_caches(UpperCamelCase_							,							UpperCamelCase_      )
																				elif isinstance(UpperCamelCase_							,							UpperCamelCase_      ):
																											SCREAMING_SNAKE_CASE__     :Optional[int]						    = [v for _, v in sorted(aa.items()							,							key=lambda UpperCamelCase_      : x[0]      )]
																											SCREAMING_SNAKE_CASE__     :Tuple						    = [v for _, v in sorted(aa.items()							,							key=lambda UpperCamelCase_      : x[0]      )]
																											consistent &= self._compare_arg_caches(UpperCamelCase_							,							UpperCamelCase_      )
																				else:
																											consistent &= aa == aa

													return consistent




						def      __lowerCamelCase       (    self :       Union[str, Any]							,							UpperCamelCase_ :       Callable							,							UpperCamelCase_ :       tuple							,							UpperCamelCase_ :       int							,							)     ->				int:
													SCREAMING_SNAKE_CASE__     :Dict						    = True
													SCREAMING_SNAKE_CASE__     :tuple						    = tree_map(lambda UpperCamelCase_      : a.shape if isinstance(UpperCamelCase_							,							torch.Tensor      ) else a							,							UpperCamelCase_							,							UpperCamelCase_      )
													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(UpperCamelCase_      )
																				SCREAMING_SNAKE_CASE__     :Optional[Any]						    = self._compare_arg_caches(self.cached_arg_data							,							UpperCamelCase_      )
													else:
																				# Otherwise, we can reuse the precomputed value
																				SCREAMING_SNAKE_CASE__     :Optional[Any]						    = False

													if not consistent:
																				SCREAMING_SNAKE_CASE__     :Optional[int]						    = self._determine_favorable_chunk_size(
																				    UpperCamelCase_							,							UpperCamelCase_							,							UpperCamelCase_							,							)
																				SCREAMING_SNAKE_CASE__     :int						    = arg_data

													assert self.cached_chunk_size is not None

													return self.cached_chunk_size






