
"""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    ):






  """simple docstring"""



  _lowercase:      List[str]   			=  []
  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			,					_UpperCamelCase    ):






  """simple docstring"""



  _lowercase:      Tuple   			=  []
  for d in reversed(_UpperCamelCase    ):
    idx.append(flat_idx % d    )
    _lowercase:      List[str]   			=  flat_idx // d

  return tuple(reversed(_UpperCamelCase    )    )




@torch.jit.ignore
def 			_lowerCAmelCase		(		_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase = None			,					_UpperCamelCase = None			,					):






  """simple docstring"""



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

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

  _lowercase:      List[Tuple[slice, ...]]   			=  []
  _lowercase:      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

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

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

    _lowercase:      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()    )
    _lowercase:      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			,					_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase    ):






  """simple docstring"""



  _lowercase:      List[str]   			=  t.shape[:no_batch_dims]
  _lowercase:      List[Any]   			=  list(_flat_idx_to_idx(_UpperCamelCase			,					_UpperCamelCase    )    )
  # _get_minimal_slice_set is inclusive
  _lowercase:      List[Any]   			=  list(_flat_idx_to_idx(flat_end - 1			,					_UpperCamelCase    )    )

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

  _lowercase:      Any   			=  [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			,					_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase			,					_UpperCamelCase = False			,					_UpperCamelCase = None			,					_UpperCamelCase = False			,					):






  """simple docstring"""



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

  _lowercase:      Any   			=  [shape[:no_batch_dims] for shape in _fetch_dims(_UpperCamelCase    )]
  _lowercase:      Tuple   			=  tuple([max(_UpperCamelCase    ) for s in zip(*_UpperCamelCase    )]    )

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

  _lowercase:      Dict[str, Any]   			=  tensor_tree_map(_prep_inputs			,					_UpperCamelCase    )
  _lowercase:      Union[str, Any]   			=  None
  if _out is not None:
    _lowercase:      List[str]   			=  tensor_tree_map(lambda _UpperCamelCase    : t.view([-1] + list(t.shape[no_batch_dims:]    )    )			,					_out    )

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

  _lowercase:      Union[str, Any]   			=  flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

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

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

    _lowercase:      Dict[str, Any]   			=  tensor_tree_map(_UpperCamelCase			,					_UpperCamelCase    )

    # Run the layer on the chunk
    _lowercase:      int   			=  layer(**_UpperCamelCase    )

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

    i += chunk_size

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

  return out







class       __magic_name__							:
       def __init__(  self			,					A_ = 512			,					)  ->						Tuple:





         """simple docstring"""




         _lowercase:      List[Any]   			=  max_chunk_size
         _lowercase:      Optional[int]   			=  None
         _lowercase:      Optional[tuple]   			=  None
       def 				lowercase_				(  self			,					A_			,					A_			,					A_							)  ->						int:





         """simple docstring"""




         logging.info('''Tuning chunk size...'''							)

         if min_chunk_size >= self.max_chunk_size:
           return min_chunk_size

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

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

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

         return candidates[min_viable_chunk_size_index]
       def 				lowercase_				(  self			,					A_			,					A_							)  ->						Any:





         """simple docstring"""




         _lowercase:      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__							):
             _lowercase:      int   			=  [v for _, v in sorted(aa.items()			,					key=lambda A_							: x[0]							)]
             _lowercase:      str   			=  [v for _, v in sorted(aa.items()			,					key=lambda A_							: x[0]							)]
             consistent &= self._compare_arg_caches(lowerCAmelCase__			,					lowerCAmelCase__							)
           else:
             consistent &= aa == aa

         return consistent



       def 				lowercase_				(  self			,					A_			,					A_			,					A_			,					)  ->						Any:





         """simple docstring"""




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

         if not consistent:
           _lowercase:      Optional[Any]   			=  self._determine_favorable_chunk_size(
               lowerCAmelCase__			,					lowerCAmelCase__			,					lowerCAmelCase__			,					)
           _lowercase:      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]					)




