



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

							# 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(A	:					List[bool]					) -> None:
														SCREAMING_SNAKE_CASE       :   Optional[Any]         =      True
														for i in range(len(A					)					):
																					SCREAMING_SNAKE_CASE       :   int         =      -1 * (i + 1)
																					l[reversed_idx] &= tally
																					SCREAMING_SNAKE_CASE       :   Tuple         =      l[reversed_idx]

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

							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(A       ,		A					):
														if s == e:
																					path_list.append(slice(A       ,		s + 1					)					)
														else:
																					break

							SCREAMING_SNAKE_CASE       :   Tuple[slice, ...]         =      tuple(A					)
							SCREAMING_SNAKE_CASE       :   Dict         =      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

														SCREAMING_SNAKE_CASE       :   List[str]         =      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

														SCREAMING_SNAKE_CASE       :   Any         =      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()					)
														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        UpperCAmelCase (    A	:					torch.Tensor       ,		A	:					int       ,		A	:					int       ,		A	:					int					):
							SCREAMING_SNAKE_CASE       :   List[str]         =      t.shape[:no_batch_dims]
							SCREAMING_SNAKE_CASE       :   List[Any]         =      list(_flat_idx_to_idx(A       ,		A					)					)
							# _get_minimal_slice_set is inclusive
							SCREAMING_SNAKE_CASE       :   List[Any]         =      list(_flat_idx_to_idx(flat_end - 1       ,		A					)					)

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

							SCREAMING_SNAKE_CASE       :   Any         =      [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       ,		):
							if not (len(A					) > 0):
														raise ValueError('''Must provide at least one input'''					)

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

							SCREAMING_SNAKE_CASE       :   Dict[str, Any]         =      tensor_tree_map(_prep_inputs       ,		A					)
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      None
							if _out is not None:
														SCREAMING_SNAKE_CASE       :   List[str]         =      tensor_tree_map(lambda A					: 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       :   Union[str, Any]         =      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

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

														SCREAMING_SNAKE_CASE       :   Dict[str, Any]         =      tensor_tree_map(A       ,		A					)

														# Run the layer on the chunk
														SCREAMING_SNAKE_CASE       :   int         =      layer(**A					)

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

														i += chunk_size

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

							return out






class 						lowerCamelCase_				:







						def __init__(	self :   Dict					,     lowerCAmelCase__ :   int = 5_12					,     ):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      max_chunk_size
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      None
													SCREAMING_SNAKE_CASE       :   Optional[tuple]         =      None







						def        __lowercase     (	self :   Any					,     lowerCAmelCase__ :   Callable					,     lowerCAmelCase__ :   tuple					,     lowerCAmelCase__ :   int					):

													"""simple docstring"""
													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[int]         =      [c for c in candidates if c > min_chunk_size]
													SCREAMING_SNAKE_CASE       :   Any         =      [min_chunk_size] + candidates
													candidates[-1] += 4

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

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

													return candidates[min_viable_chunk_size_index]







						def        __lowercase     (	self :   str					,     lowerCAmelCase__ :   Iterable					,     lowerCAmelCase__ :   Iterable					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      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__					):
																											SCREAMING_SNAKE_CASE       :   int         =      [v for _, v in sorted(aa.items()					,     key=lambda lowerCAmelCase__					: x[0]					)]
																											SCREAMING_SNAKE_CASE       :   str         =      [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        __lowercase     (	self :   List[str]					,     lowerCAmelCase__ :   Callable					,     lowerCAmelCase__ :   tuple					,     lowerCAmelCase__ :   int					,     ):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   int         =      True
													SCREAMING_SNAKE_CASE       :   tuple         =      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__					)
																				SCREAMING_SNAKE_CASE       :   Optional[int]         =      self._compare_arg_caches(self.cached_arg_data					,     lowerCAmelCase__					)
													else:
																				# Otherwise, we can reuse the precomputed value
																				SCREAMING_SNAKE_CASE       :   int         =      False

													if not consistent:
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      self._determine_favorable_chunk_size(
																				    lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     )
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      arg_data

													assert self.cached_chunk_size is not None

													return self.cached_chunk_size









'''simple docstring'''



import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.eager import context
							from tensorflow.python.framework import ops

							from transformers import GradientAccumulator, create_optimizer
@require_tf
class 						lowerCamelCase_				(       unittest.TestCase      ):







						def        __lowercase     (	self :   List[Any]					,     lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   Dict					,     lowerCAmelCase__ :   int					):

													"""simple docstring"""
													self.assertEqual(len(lowerCAmelCase__					)					,     len(lowerCAmelCase__					)					)
													for a, b in zip(lowerCAmelCase__					,     lowerCAmelCase__					):
																				self.assertAlmostEqual(lowerCAmelCase__					,     lowerCAmelCase__					,     delta=lowerCAmelCase__					)







						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      GradientAccumulator()
													accumulator([tf.constant([1.0, 2.0]					)]					)
													accumulator([tf.constant([-2.0, 1.0]					)]					)
													accumulator([tf.constant([-1.0, 2.0]					)]					)
													with self.assertRaises(lowerCAmelCase__					):
																				accumulator([tf.constant([1.0, 1.0]					), tf.constant([2.0, 2.0]					)]					)
													self.assertEqual(accumulator.step					,     3					)
													self.assertEqual(len(accumulator.gradients					)					,     1					)
													self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()					,     [-2.0, 5.0]					,     tol=1e-2					)
													accumulator.reset()
													self.assertEqual(accumulator.step					,     0					)
													self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()					,     [0.0, 0.0]					,     tol=1e-2					)







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      None
													ops.enable_eager_execution_internal()
													SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.config.list_physical_devices('''CPU'''					)
													if len(lowerCAmelCase__					) == 1:
																				tf.config.set_logical_device_configuration(
																				    physical_devices[0]					,     [tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]					)
													SCREAMING_SNAKE_CASE       :   Tuple         =      tf.config.list_logical_devices(device_type='''CPU'''					)
													SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.distribute.MirroredStrategy(devices=devices[:2]					)

													with strategy.scope():
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      GradientAccumulator()
																				SCREAMING_SNAKE_CASE       :   Dict         =      tf.Variable([4.0, 3.0]					)
																				SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Optional[int]         =      create_optimizer(5e-5					,     10					,     5					)
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      tf.Variable([0.0, 0.0]					,     trainable=lowerCAmelCase__					)

													def accumulate_on_replica(lowerCAmelCase__ :   List[str]					):
																				accumulator([gradient]					)

													def apply_on_replica():
																				optimizer.apply_gradients(list(zip(accumulator.gradients					,     [variable]					)					)					)

													@tf.function
													def accumulate(lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   List[Any]					):
																				with strategy.scope():
																											SCREAMING_SNAKE_CASE       :   Optional[Any]         =      strategy.experimental_local_results(lowerCAmelCase__					)
																											local_variables[0].assign(lowerCAmelCase__					)
																											local_variables[1].assign(lowerCAmelCase__					)
																											strategy.run(lowerCAmelCase__					,     args=(gradient_placeholder,)					)

													@tf.function
													def apply_grad():
																				with strategy.scope():
																											strategy.run(lowerCAmelCase__					)

													def _check_local_values(lowerCAmelCase__ :   Dict					,     lowerCAmelCase__ :   int					):
																				SCREAMING_SNAKE_CASE       :   Tuple         =      strategy.experimental_local_results(accumulator._gradients[0]					)
																				self.assertListAlmostEqual(values[0].value()					,     lowerCAmelCase__					,     tol=1e-2					)
																				self.assertListAlmostEqual(values[1].value()					,     lowerCAmelCase__					,     tol=1e-2					)

													accumulate([1.0, 2.0]					,     [-1.0, 1.0]					)
													accumulate([3.0, -1.0]					,     [-1.0, -1.0]					)
													accumulate([-2.0, 2.0]					,     [3.0, -2.0]					)
													self.assertEqual(accumulator.step					,     3					)
													_check_local_values([2.0, 3.0]					,     [1.0, -2.0]					)
													apply_grad()
													self.assertListAlmostEqual(variable.value()					,     [4.0, 3.0]					,     tol=1e-2					)
													accumulator.reset()
													self.assertEqual(accumulator.step					,     0					)
													_check_local_values([0.0, 0.0]					,     [0.0, 0.0]					)




