



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





						"""simple docstring"""





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

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





						"""simple docstring"""





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

						return tuple(reversed(_A    )    )
@torch.jit.ignore
def  UpperCamelCase						(       _A			:				Sequence[int]			,							_A			:				Sequence[int]			,							_A			:				Sequence[int]			,							_A			:				Optional[Sequence[bool]] = None			,							_A			:				Optional[Sequence[bool]] = None			,							)->       List[Tuple[slice, ...]]:





						"""simple docstring"""





						def reduce_edge_list(_A			:				List[bool]    ) -> None:
												A__							=       True
												for i in range(len(_A    )    ):
																		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(_A    )
						if end_edges is None:
												A__							=       [e == (d - 1) for e, d in zip(_A			,							_A    )]
												reduce_edge_list(_A    )

						# Base cases. Either start/end are empty and we're done, or the final,
						# one-dimensional tensor can be simply sliced
						if len(_A    ) == 0:
												return [()]
						elif len(_A    ) == 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(_A			,							_A    ):
												if s == e:
																		path_list.append(slice(_A			,							s + 1    )    )
												else:
																		break

						A__							=       tuple(_A    )
						A__							=       len(_A    )

						# start == end, and we're done
						if divergence_idx == len(_A    ):
												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(_A			,							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(_A			,							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						(       _A			:				torch.Tensor			,							_A			:				int			,							_A			:				int			,							_A			:				int    )->       torch.Tensor:





						"""simple docstring"""





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

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

						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						(       _A			:				Callable			,							_A			:				Dict[str, Any]			,							_A			:				int			,							_A			:				int			,							_A			:				bool = False			,							_A			:				Any = None			,							_A			:				bool = False			,							)->       Any:





						"""simple docstring"""





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

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

						def _prep_inputs(_A			:				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			,							_A    )
						A__							=       None
						if _out is not None:
												A__							=       tensor_tree_map(lambda _A    : 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(_A			:				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(_A    ):
												# Chunk the input
												if not low_mem:
																		A__							=       _select_chunk
												else:
																		A__							=       partial(
																		    _chunk_slice			,							flat_start=_A			,							flat_end=min(_A			,							i + chunk_size    )			,							no_batch_dims=len(_A    )			,							)

												A__							=       tensor_tree_map(_A			,							_A    )

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

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

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

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

																		assign(_A			,							_A    )
												elif isinstance(_A			,							_A    ):
																		for xa, xa in zip(_A			,							_A    ):
																								if _add_into_out:
																														xa[i : i + chunk_size] += xa
																								else:
																														A__							=       xa
												elif isinstance(_A			,							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 _A    : t.view(orig_batch_dims + t.shape[1:]    )			,							_A    )

						return out






class      UpperCamelCase    :



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



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							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(UpperCAmelCase__		) -> bool:
													try:
																			with torch.no_grad():
																									fn(*UpperCAmelCase__			,	chunk_size=UpperCAmelCase__		)
																			return True
													except RuntimeError:
																			return False

							A__							=       0
							A__							=       len(UpperCAmelCase__		) - 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(UpperCAmelCase__		) - 1) // 2

							return candidates[min_viable_chunk_size_index]



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       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__		):
																			A__							=       [v for _, v in sorted(aa.items()			,	key=lambda UpperCAmelCase__		: x[0]		)]
																			A__							=       [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 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	):
							A__							=       True
							A__							=       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__		)
													A__							=       self._compare_arg_caches(self.cached_arg_data			,	UpperCAmelCase__		)
							else:
													# Otherwise, we can reuse the precomputed value
													A__							=       False

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

							assert self.cached_chunk_size is not None

							return self.cached_chunk_size










def  UpperCamelCase						(       _A			:				str    )->       str:





						"""simple docstring"""





						A__							=       ""
						for ch in key:
												if ch == " " or ch not in key_no_dups and ch.isalpha():
																		key_no_dups += ch
						return key_no_dups
def  UpperCamelCase						(       _A			:				str    )->       dict[str, str]:





						"""simple docstring"""





						A__							=       [chr(i + 65    ) for i in range(26    )]
						# Remove duplicate characters from key
						A__							=       remove_duplicates(key.upper()    )
						A__							=       len(_A    )
						# First fill cipher with key characters
						A__							=       {alphabet[i]: char for i, char in enumerate(_A    )}
						# Then map remaining characters in alphabet to
						# the alphabet from the beginning
						for i in range(len(_A    )			,							26    ):
												A__							=       alphabet[i - offset]
												# Ensure we are not mapping letters to letters previously mapped
												while char in key:
																		offset -= 1
																		A__							=       alphabet[i - offset]
												A__							=       char
						return cipher_alphabet
def  UpperCamelCase						(       _A			:				str			,							_A			:				dict[str, str]    )->       str:





						"""simple docstring"""





						return "".join(cipher_map.get(_A			,							_A    ) for ch in message.upper()    )
def  UpperCamelCase						(       _A			:				str			,							_A			:				dict[str, str]    )->       str:





						"""simple docstring"""





						A__							=       {v: k for k, v in cipher_map.items()}
						return "".join(rev_cipher_map.get(_A			,							_A    ) for ch in message.upper()    )
def  UpperCamelCase						(       )->       None:





						"""simple docstring"""





						A__							=       input("Enter message to encode or decode: "    ).strip()
						A__							=       input("Enter keyword: "    ).strip()
						A__							=       input("Encipher or decipher? E/D:"    ).strip()[0].lower()
						try:
												A__							=       {"e": encipher, "d": decipher}[option]
						except KeyError:
												raise KeyError("invalid input option"    )
						A__							=       create_cipher_map(_A    )
						print(func(_A			,							_A    )    )


if __name__ == "__main__":
						import doctest

						doctest.testmod()
						main()





