


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_   (	_a	):




		"""simple docstring"""






		lowerCAmelCase__       :      str   					= []
		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 					lowerCamelCase_   (	_a     ,       _a	):




		"""simple docstring"""






		lowerCAmelCase__       :      str   					= []
		for d in reversed(_a	):
				idx.append(flat_idx % d	)
				lowerCAmelCase__       :      Optional[Any]   					= flat_idx // d

		return tuple(reversed(_a	)	)






@torch.jit.ignore
def 					lowerCamelCase_   (	_a     ,       _a     ,       _a     ,       _a = None     ,       _a = None     ,       ):




		"""simple docstring"""






		def reduce_edge_list(_a	) -> None:
				lowerCAmelCase__       :      Union[str, Any]   					= True
				for i in range(len(_a	)	):
						lowerCAmelCase__       :      Optional[int]   					= -1 * (i + 1)
						l[reversed_idx] &= tally
						lowerCAmelCase__       :      List[Any]   					= l[reversed_idx]

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

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

		# 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

		lowerCAmelCase__       :      Tuple[slice, ...]   					= tuple(_a	)
		lowerCAmelCase__       :      List[Any]   					= 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

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

				lowerCAmelCase__       :      int   					= 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()	)
				lowerCAmelCase__       :      List[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_   (	_a     ,       _a     ,       _a     ,       _a	):




		"""simple docstring"""






		lowerCAmelCase__       :      Optional[int]   					= t.shape[:no_batch_dims]
		lowerCAmelCase__       :      Tuple   					= list(_flat_idx_to_idx(_a     ,       _a	)	)
		# _get_minimal_slice_set is inclusive
		lowerCAmelCase__       :      List[str]   					= list(_flat_idx_to_idx(flat_end - 1     ,       _a	)	)

		# Get an ordered list of slices to perform
		lowerCAmelCase__       :      Tuple   					= _get_minimal_slice_set(
		    _a     ,       _a     ,       _a     ,       )

		lowerCAmelCase__       :      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_   (	_a     ,       _a     ,       _a     ,       _a     ,       _a = False     ,       _a = None     ,       _a = False     ,       ):




		"""simple docstring"""






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

		lowerCAmelCase__       :      int   					= [shape[:no_batch_dims] for shape in _fetch_dims(_a	)]
		lowerCAmelCase__       :      Tuple   					= tuple([max(_a	) for s in zip(*_a	)]	)

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

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

		lowerCAmelCase__       :      str   					= 1
		for d in orig_batch_dims:
				flat_batch_dim *= d

		lowerCAmelCase__       :      Optional[Any]   					= flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

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

		lowerCAmelCase__       :      List[Any]   					= 0
		lowerCAmelCase__       :      str   					= prepped_outputs
		for _ in range(_a	):
				# Chunk the input
				if not low_mem:
						lowerCAmelCase__       :      str   					= _select_chunk
				else:
						lowerCAmelCase__       :      Optional[int]   					= partial(
						    _chunk_slice     ,       flat_start=_a     ,       flat_end=min(_a     ,       i + chunk_size	)     ,       no_batch_dims=len(_a	)     ,       )

				lowerCAmelCase__       :      Dict[str, Any]   					= tensor_tree_map(_a     ,       _a	)

				# Run the layer on the chunk
				lowerCAmelCase__       :      str   					= layer(**_a	)

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

				i += chunk_size

		lowerCAmelCase__       :      Dict   					= tensor_tree_map(lambda _a	: t.view(orig_batch_dims + t.shape[1:]	)     ,       _a	)

		return out





class  _a							:




	def __init__( self						:		Any		, _SCREAMING_SNAKE_CASE						:		int = 512		, )->	Tuple:
			lowerCAmelCase__       :      Dict   					= max_chunk_size
			lowerCAmelCase__       :      Optional[int]   					= None
			lowerCAmelCase__       :      Optional[tuple]   					= None




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

			if min_chunk_size >= self.max_chunk_size:
					return min_chunk_size

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

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

			lowerCAmelCase__       :      Tuple   					= 0
			lowerCAmelCase__       :      Dict   					= len(_SCREAMING_SNAKE_CASE		) - 1
			while i > min_viable_chunk_size_index:
					lowerCAmelCase__       :      int   					= test_chunk_size(candidates[i]		)
					if not viable:
							lowerCAmelCase__       :      Dict   					= (min_viable_chunk_size_index + i) // 2
					else:
							lowerCAmelCase__       :      Optional[int]   					= i
							lowerCAmelCase__       :      Optional[int]   					= (i + len(_SCREAMING_SNAKE_CASE		) - 1) // 2

			return candidates[min_viable_chunk_size_index]




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

			return consistent




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

			if not consistent:
					lowerCAmelCase__       :      str   					= self._determine_favorable_chunk_size(
					    _SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		, _SCREAMING_SNAKE_CASE		, )
					lowerCAmelCase__       :      Dict   					= arg_data

			assert self.cached_chunk_size is not None

			return self.cached_chunk_size








import timeit

import numpy as np

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features.features import _ArrayXD






def 					lowerCamelCase_   (	_a	):




		"""simple docstring"""






		def wrapper(*_a     ,       **_a	):
				lowerCAmelCase__       :      List[str]   					= timeit.default_timer()
				lowerCAmelCase__       :      List[Any]   					= func(*_a     ,       **_a	)
				lowerCAmelCase__       :      Any   					= timeit.default_timer() - starttime
				return delta

		lowerCAmelCase__       :      Any   					= func.__name__

		return wrapper






def 					lowerCamelCase_   (	_a     ,       _a=100     ,       _a=None	):




		"""simple docstring"""






		lowerCAmelCase__       :      str   					= []
		lowerCAmelCase__       :      str   					= seq_shapes or {}
		for i in range(_a	):
				lowerCAmelCase__       :      List[str]   					= {}
				for col_id, (k, v) in enumerate(features.items()	):
						if isinstance(_a     ,       _ArrayXD	):
								lowerCAmelCase__       :      List[str]   					= np.random.rand(*v.shape	).astype(v.dtype	)
						elif isinstance(_a     ,       datasets.Value	):
								if v.dtype == "string":
										lowerCAmelCase__       :      Dict   					= '''The small grey turtle was surprisingly fast when challenged.'''
								else:
										lowerCAmelCase__       :      Any   					= np.random.randint(10     ,       size=1	).astype(v.dtype	).item()
						elif isinstance(_a     ,       datasets.Sequence	):
								while isinstance(_a     ,       datasets.Sequence	):
										lowerCAmelCase__       :      Optional[int]   					= v.feature
								lowerCAmelCase__       :      str   					= seq_shapes[k]
								lowerCAmelCase__       :      Any   					= np.random.rand(*_a	).astype(v.dtype	)
						lowerCAmelCase__       :      int   					= data

				dummy_data.append((i, example)	)

		return dummy_data






def 					lowerCamelCase_   (	_a     ,       _a     ,       _a=100     ,       _a=None	):




		"""simple docstring"""






		lowerCAmelCase__       :      Optional[Any]   					= generate_examples(_a     ,       num_examples=_a     ,       seq_shapes=_a	)

		with ArrowWriter(features=_a     ,       path=_a	) as writer:
				for key, record in dummy_data:
						lowerCAmelCase__       :      Optional[int]   					= features.encode_example(_a	)
						writer.write(_a	)

				lowerCAmelCase__					,     lowerCAmelCase__       :      List[str]   					= writer.finalize()

		if not num_final_examples == num_examples:
				raise ValueError(
				    f'Error writing the dataset, wrote {num_final_examples} examples but should have written {num_examples}.'	)

		lowerCAmelCase__       :      List[Any]   					= datasets.Dataset.from_file(filename=_a     ,       info=datasets.DatasetInfo(features=_a	)	)

		return dataset




