



'''simple docstring'''





import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


lowercase_			=  namedtuple("ExampleNamedTuple", "a b c")






class __A      (  unittest.TestCase						):



 '''simple docstring'''







 def  a__  (self	)   ->					List[Any]:


      """simple docstring"""



      _a     =  torch.randn(5			,       2	)
      _a     =  torch.device('''cuda'''	) if torch.cuda.is_available() else torch.device('''cpu'''	)

      _a     =  send_to_device(A			,       A	)
      self.assertTrue(torch.equal(resulta.cpu()			,       A	)	)

      _a     =  send_to_device((tensor, [tensor, tensor], 1)			,       A	)
      self.assertIsInstance(A			,       A	)
      self.assertTrue(torch.equal(resulta[0].cpu()			,       A	)	)
      self.assertIsInstance(resulta[1]			,       A	)
      self.assertTrue(torch.equal(resulta[1][0].cpu()			,       A	)	)
      self.assertTrue(torch.equal(resulta[1][1].cpu()			,       A	)	)
      self.assertEqual(resulta[2]			,       1	)

      _a     =  send_to_device({'''a''': tensor, '''b''': [tensor, tensor], '''c''': 1}			,       A	)
      self.assertIsInstance(A			,       A	)
      self.assertTrue(torch.equal(resulta['''a'''].cpu()			,       A	)	)
      self.assertIsInstance(resulta['''b''']			,       A	)
      self.assertTrue(torch.equal(resulta['''b'''][0].cpu()			,       A	)	)
      self.assertTrue(torch.equal(resulta['''b'''][1].cpu()			,       A	)	)
      self.assertEqual(resulta['''c''']			,       1	)

      _a     =  send_to_device(ExampleNamedTuple(a=A			,       b=[tensor, tensor]			,       c=1	)			,       A	)
      self.assertIsInstance(A			,       A	)
      self.assertTrue(torch.equal(resulta.a.cpu()			,       A	)	)
      self.assertIsInstance(resulta.b			,       A	)
      self.assertTrue(torch.equal(resulta.b[0].cpu()			,       A	)	)
      self.assertTrue(torch.equal(resulta.b[1].cpu()			,       A	)	)
      self.assertEqual(resulta.c			,       1	)

      _a     =  send_to_device(UserDict({'''a''': tensor, '''b''': [tensor, tensor], '''c''': 1}	)			,       A	)
      self.assertIsInstance(A			,       A	)
      self.assertTrue(torch.equal(resulta['''a'''].cpu()			,       A	)	)
      self.assertIsInstance(resulta['''b''']			,       A	)
      self.assertTrue(torch.equal(resulta['''b'''][0].cpu()			,       A	)	)
      self.assertTrue(torch.equal(resulta['''b'''][1].cpu()			,       A	)	)
      self.assertEqual(resulta['''c''']			,       1	)



 def  a__  (self	)   ->					int:


      """simple docstring"""



      with self.assertRaises(A	) as cm:
           _a     =  recursively_apply(torch.tensor			,       (torch.tensor(1	), 1)			,       error_on_other_type=A	)
      self.assertEqual(
          str(cm.exception	)			,       '''Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'''			,       )



 def  a__  (self	)   ->					Optional[int]:


      """simple docstring"""



      _a     =  torch.tensor([1, 2, 3, 4, 5]	)
      self.assertEqual(listify(A	)			,       [1, 2, 3, 4, 5]	)

      _a     =  torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]	)
      self.assertEqual(listify(A	)			,       [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]	)

      _a     =  torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]	)
      self.assertEqual(
          listify(A	)			,       [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]	)



 def  a__  (self	)   ->					Optional[Any]:


      """simple docstring"""



      with patch_environment(aa=1			,       BB=2	):
           self.assertEqual(os.environ.get('''AA'''	)			,       '''1'''	)
           self.assertEqual(os.environ.get('''BB'''	)			,       '''2'''	)

      self.assertNotIn('''AA'''			,       os.environ	)
      self.assertNotIn('''BB'''			,       os.environ	)



 def  a__  (self	)   ->					Union[str, Any]:


      """simple docstring"""



      _a     =  RegressionModel()
      _a     =  model.forward
      _a     =  convert_outputs_to_fpaa(model.forward	)
      _a     =  extract_model_from_parallel(A			,       keep_fpaa_wrapper=A	)
      _a     =  pickle.dumps(A	)



 @require_cuda
 def  a__  (self	)   ->					List[str]:


      """simple docstring"""



      _a     =  RegressionModel()
      _a     =  model.forward
      _a     =  torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
      _a     =  convert_outputs_to_fpaa(model.forward	)
      _a     =  extract_model_from_parallel(A			,       keep_fpaa_wrapper=A	)
      _a     =  pickle.dumps(A	)



 @require_cuda
 @require_torch_min_version(version='''2.0'''	)
 def  a__  (self	)   ->					Dict:


      """simple docstring"""



      _a     =  RegressionModel()
      _a     =  model.forward
      _a     =  torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
      _a     =  convert_outputs_to_fpaa(model.forward	)
      _a     =  torch.compile(model.forward			,       backend='''inductor'''	)
      _a     =  torch.randn(4			,       10	).cuda()
      _a     =  model(A	)



 def  a__  (self	)   ->					List[str]:


      """simple docstring"""



      _a     =  RegressionModel()
      # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
      _a     =  torch.nn.parallel.DataParallel(A	)
      _a     =  extract_model_from_parallel(A	)

      self.assertEqual(A			,       A	)



 @require_torch_min_version(version='''2.0'''	)
 def  a__  (self	)   ->					Tuple:


      """simple docstring"""



      _a     =  RegressionModel()
      _a     =  torch.compile(A	)

      # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
      _a     =  torch.nn.parallel.DataParallel(A	)
      _a     =  torch.compile(A	)
      _a     =  extract_model_from_parallel(A	)

      self.assertEqual(compiled_model._orig_mod			,       compiled_model_unwrapped._orig_mod	)



 def  a__  (self	)   ->					int:


      """simple docstring"""



      self.assertEqual(find_device([1, '''a''', torch.tensor([1, 2, 3]	)]	)			,       torch.device('''cpu'''	)	)
      self.assertEqual(find_device({'''a''': 1, '''b''': torch.tensor([1, 2, 3]	)}	)			,       torch.device('''cpu'''	)	)
      self.assertIsNone(find_device([1, '''a''']	)	)








'''simple docstring'''





import warnings
from collections.abc import Callable
from typing import Any

import numpy
from matplotlib import pyplot

lowercase_			=  0.25 + 0.0J
lowercase_			=  -0.4 + 0.6J
lowercase_			=  -0.1 + 0.651J
lowercase_			=  -2.0
lowercase_			=  56
lowercase_			=  2.0
lowercase_			=  666




def lowerCAmelCase  (__A  ,     __A):



     """simple docstring"""



     return numpy.exp(__A) + c_parameter




def lowerCAmelCase  (__A  ,     __A):



     """simple docstring"""



     return z_values * z_values + c_parameter




def lowerCAmelCase  (__A  ,     __A):



     """simple docstring"""



     _a     =  numpy.linspace(-window_size  ,     __A  ,     __A)
     _a     =  x.reshape((nb_pixels, 1))
     _a     =  numpy.linspace(-window_size  ,     __A  ,     __A)
     _a     =  y.reshape((1, nb_pixels))
     return x + 1.0J * y




def lowerCAmelCase  (__A  ,     __A  ,     __A  ,     __A  ,     __A = None  ,     ):



     """simple docstring"""



     _a     =  z_a.astype('''complex64''')
     for _ in range(__A):
          _a     =  eval_function(__A  ,     __A)
          if infinity is not None:
               numpy.nan_to_num(__A  ,     copy=__A  ,     nan=__A)
               _a     =  infinity
     return z_n




def lowerCAmelCase  (__A  ,     __A  ,     __A  ,     __A  ,     ):



     """simple docstring"""



     _a     =  (abs(__A)).transpose()
     _a     =  abs_z_final[::-1, :]
     pyplot.matshow(abs_z_final < escape_radius)
     pyplot.title(F'''Julia set of ${function_label}$, $c={function_params}$''')
     pyplot.show()




def lowerCAmelCase  ():



     """simple docstring"""



     warnings.filterwarnings(
         '''ignore'''  ,     category=__A  ,     message='''overflow encountered in multiply''')
     warnings.filterwarnings(
         '''ignore'''  ,     category=__A  ,     message='''invalid value encountered in multiply'''  ,     )
     warnings.filterwarnings(
         '''ignore'''  ,     category=__A  ,     message='''overflow encountered in absolute''')
     warnings.filterwarnings(
         '''ignore'''  ,     category=__A  ,     message='''overflow encountered in exp''')


if __name__ == "__main__":
  lowercase_			=  prepare_grid(window_size, nb_pixels)

  ignore_overflow_warnings()  # See file header for explanations

  lowercase_			=  24
  lowercase_			=  2 * abs(c_cauliflower) + 1
  lowercase_			=  iterate_function(
      eval_quadratic_polynomial,
      c_cauliflower,
      nb_iterations,
      z_a,
      infinity=1.1 * escape_radius,
  )
  show_results("z^2+c", c_cauliflower, escape_radius, z_final)

  lowercase_			=  64
  lowercase_			=  2 * abs(c_polynomial_a) + 1
  lowercase_			=  iterate_function(
      eval_quadratic_polynomial,
      c_polynomial_a,
      nb_iterations,
      z_a,
      infinity=1.1 * escape_radius,
  )
  show_results("z^2+c", c_polynomial_a, escape_radius, z_final)

  lowercase_			=  161
  lowercase_			=  2 * abs(c_polynomial_a) + 1
  lowercase_			=  iterate_function(
      eval_quadratic_polynomial,
      c_polynomial_a,
      nb_iterations,
      z_a,
      infinity=1.1 * escape_radius,
  )
  show_results("z^2+c", c_polynomial_a, escape_radius, z_final)

  lowercase_			=  12
  lowercase_			=  1_0000.0
  lowercase_			=  iterate_function(
      eval_exponential,
      c_exponential,
      nb_iterations,
      z_a + 2,
      infinity=1.0e10,
  )
  show_results("e^z+c", c_exponential, escape_radius, z_final)



