

import os
import unittest
from tempfile import TemporaryDirectory

import torch
import torch.nn as nn

from accelerate.utils import (
    OffloadedWeightsLoader,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
)



class       __snake_case  (    nn.Module       ):

  def __init__(					self):
      super().__init__()
      SCREAMING_SNAKE_CASE_									=						nn.Linear(3						,		4)
      SCREAMING_SNAKE_CASE_									=						nn.BatchNormad(4)
      SCREAMING_SNAKE_CASE_									=						nn.Linear(4						,		5)


  def 						lowerCAmelCase__     (					self						,		_A):
      return self.lineara(self.batchnorm(self.lineara(_A)))






class       __snake_case  (    unittest.TestCase       ):

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						ModelForTest()
      with TemporaryDirectory() as tmp_dir:
          offload_state_dict(_A						,		model.state_dict())
          SCREAMING_SNAKE_CASE_									=						os.path.join(_A						,		'index.json')
          self.assertTrue(os.path.isfile(_A))
          # TODO: add tests on what is inside the index

          for key in ["linear1.weight", "linear1.bias", "linear2.weight", "linear2.bias"]:
              SCREAMING_SNAKE_CASE_									=						os.path.join(_A						,		f"""{key}.dat""")
              self.assertTrue(os.path.isfile(_A))
              # TODO: add tests on the fact weights are properly loaded

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						[torch.floataa, torch.floataa, torch.bfloataa]

      for dtype in dtypes:
          SCREAMING_SNAKE_CASE_									=						torch.randn(2						,		3						,		dtype=_A)
          with TemporaryDirectory() as tmp_dir:
              SCREAMING_SNAKE_CASE_									=						offload_weight(_A						,		'weight'						,		_A						,		{})
              SCREAMING_SNAKE_CASE_									=						os.path.join(_A						,		'weight.dat')
              self.assertTrue(os.path.isfile(_A))
              self.assertDictEqual(_A						,		{'weight': {'shape': [2, 3], 'dtype': str(_A).split('.')[1]}})

              SCREAMING_SNAKE_CASE_									=						load_offloaded_weight(_A						,		index['weight'])
              self.assertTrue(torch.equal(_A						,		_A))

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						ModelForTest()
      SCREAMING_SNAKE_CASE_									=						model.state_dict()
      SCREAMING_SNAKE_CASE_									=						{k: v for k, v in state_dict.items() if 'linear2' not in k}
      SCREAMING_SNAKE_CASE_									=						{k: v for k, v in state_dict.items() if 'linear2' in k}

      with TemporaryDirectory() as tmp_dir:
          offload_state_dict(_A						,		_A)
          SCREAMING_SNAKE_CASE_									=						OffloadedWeightsLoader(state_dict=_A						,		save_folder=_A)

          # Every key is there with the right value
          self.assertEqual(sorted(_A)						,		sorted(state_dict.keys()))
          for key, param in state_dict.items():
              self.assertTrue(torch.allclose(_A						,		weight_map[key]))

      SCREAMING_SNAKE_CASE_									=						{k: v for k, v in state_dict.items() if 'weight' in k}
      SCREAMING_SNAKE_CASE_									=						{k: v for k, v in state_dict.items() if 'weight' not in k}

      with TemporaryDirectory() as tmp_dir:
          offload_state_dict(_A						,		_A)
          SCREAMING_SNAKE_CASE_									=						OffloadedWeightsLoader(state_dict=_A						,		save_folder=_A)

          # Every key is there with the right value
          self.assertEqual(sorted(_A)						,		sorted(state_dict.keys()))
          for key, param in state_dict.items():
              self.assertTrue(torch.allclose(_A						,		weight_map[key]))

      with TemporaryDirectory() as tmp_dir:
          offload_state_dict(_A						,		_A)
          # Duplicates are removed
          SCREAMING_SNAKE_CASE_									=						OffloadedWeightsLoader(state_dict=_A						,		save_folder=_A)

          # Every key is there with the right value
          self.assertEqual(sorted(_A)						,		sorted(state_dict.keys()))
          for key, param in state_dict.items():
              self.assertTrue(torch.allclose(_A						,		weight_map[key]))


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						{'a.1': 0, 'a.10': 1, 'a.2': 2}
      SCREAMING_SNAKE_CASE_									=						extract_submodules_state_dict(_A						,		['a.1', 'a.2'])
      self.assertDictEqual(_A						,		{'a.1': 0, 'a.2': 2})

      SCREAMING_SNAKE_CASE_									=						{'a.1.a': 0, 'a.10.a': 1, 'a.2.a': 2}
      SCREAMING_SNAKE_CASE_									=						extract_submodules_state_dict(_A						,		['a.1', 'a.2'])
      self.assertDictEqual(_A						,		{'a.1.a': 0, 'a.2.a': 2})





import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command



class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Dict						 =  inspect.getfile(accelerate.test_utils       )
  __lowerCAmelCase     :	Optional[Any]						 =  os.path.sep.join(mod_file.split(os.path.sep       )[:-1] + ['scripts', 'test_cli.py']       )

  __lowerCAmelCase     :	Tuple						 =  ['accelerate', 'launch']
  __lowerCAmelCase     :	Union[str, Any]						 =  Path.home() / '.cache/huggingface/accelerate'
  __lowerCAmelCase     :	List[str]						 =  'default_config.yaml'
  __lowerCAmelCase     :	List[Any]						 =  config_folder / config_file
  __lowerCAmelCase     :	str						 =  config_folder / '_default_config.yaml'

  __lowerCAmelCase     :	Optional[int]						 =  Path('tests/test_configs'       )

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.config_path.is_file():
          cls.config_path.rename(cls.changed_path)

  @classmethod
  def 						lowerCAmelCase__     (					cls):
      if cls.changed_path.is_file():
          cls.changed_path.rename(cls.config_path)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.base_cmd
      if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
          cmd += ["--multi_gpu"]
      execute_subprocess_async(cmd + [self.test_file_path]						,		env=os.environ.copy())

  def 						lowerCAmelCase__     (					self):
      for config in sorted(self.test_config_path.glob('**/*.yaml')):
          with self.subTest(config_file=_A):
              execute_subprocess_async(
                  self.base_cmd + ['--config_file', str(_A), self.test_file_path]						,		env=os.environ.copy())


  def 						lowerCAmelCase__     (					self):
      execute_subprocess_async(['accelerate', 'test']						,		env=os.environ.copy())






class       __snake_case  (    unittest.TestCase       ):
  __lowerCAmelCase     :	Optional[Any]						 =  'test-tpu'
  __lowerCAmelCase     :	str						 =  'us-central1-a'
  __lowerCAmelCase     :	Union[str, Any]						 =  'ls'
  __lowerCAmelCase     :	Union[str, Any]						 =  ['accelerate', 'tpu-config']
  __lowerCAmelCase     :	Union[str, Any]						 =  'cd /usr/share'
  __lowerCAmelCase     :	List[Any]						 =  'tests/test_samples/test_command_file.sh'
  __lowerCAmelCase     :	Dict						 =  'Running gcloud compute tpus tpu-vm ssh'

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--command', self.command, '--tpu_zone', self.tpu_zone, '--tpu_name', self.tpu_name, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command',
              self.command,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--debug']						,		return_stdout=_A)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--command', self.command, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--command',
              self.command,
              '--command',
              'echo "Hello World"',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo \"Hello World\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + ['--config_file', 'tests/test_configs/latest.yaml', '--command_file', self.command_file, '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/0_12_0.yaml',
              '--command_file',
              self.command_file,
              '--tpu_zone',
              self.tpu_zone,
              '--tpu_name',
              self.tpu_name,
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--install_accelerate', '--debug']						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						run_command(
          self.cmd
          + [
              '--config_file',
              'tests/test_configs/latest.yaml',
              '--install_accelerate',
              '--accelerate_version',
              '12.0.0',
              '--debug',
          ]						,		return_stdout=_A						,		)
      self.assertIn(
          f"""{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo \"hello world\"; echo \"this is a second command\" --worker all"""						,		_A						,		)


