






'''simple docstring'''



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 							__SCREAMING_SNAKE_CASE       (							nn.Module     ):




      def __init__(  self      )							->Tuple:



         '''simple docstring'''
         super().__init__()
         __a				     =    nn.Linear(3 ,   4      )
         __a				     =    nn.BatchNormad(4      )
         __a				     =    nn.Linear(4 ,   5      )






      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->int:



         '''simple docstring'''
         return self.lineara(self.batchnorm(self.lineara(lowerCamelCase      )      )      )







class 							__SCREAMING_SNAKE_CASE       (							unittest.TestCase     ):




      def  __UpperCamelCase							(  self      )							->int:



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

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




      def  __UpperCamelCase							(  self      )							->Dict:



         '''simple docstring'''
         __a				     =    [torch.floataa, torch.floataa, torch.bfloataa]

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

               __a				     =    load_offloaded_weight(lowerCamelCase ,   index['weight']      )
               self.assertTrue(torch.equal(lowerCamelCase ,   lowerCamelCase      )      )




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



         '''simple docstring'''
         __a				     =    ModelForTest()
         __a				     =    model.state_dict()
         __a				     =    {k: v for k, v in state_dict.items() if 'linear2' not in k}
         __a				     =    {k: v for k, v in state_dict.items() if 'linear2' in k}

         with TemporaryDirectory() as tmp_dir:
            offload_state_dict(lowerCamelCase ,   lowerCamelCase      )
            __a				     =    OffloadedWeightsLoader(state_dict=lowerCamelCase ,   save_folder=lowerCamelCase      )

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

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

         with TemporaryDirectory() as tmp_dir:
            offload_state_dict(lowerCamelCase ,   lowerCamelCase      )
            __a				     =    OffloadedWeightsLoader(state_dict=lowerCamelCase ,   save_folder=lowerCamelCase      )

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

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

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






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



         '''simple docstring'''
         __a				     =    {'a.1': 0, 'a.10': 1, 'a.2': 2}
         __a				     =    extract_submodules_state_dict(lowerCamelCase ,   ['a.1', 'a.2']      )
         self.assertDictEqual(lowerCamelCase ,   {'a.1': 0, 'a.2': 2}      )

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







'''simple docstring'''



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	int    )  ->			bool:

   """simple docstring"""
   return number & 1 == 0


if __name__ == "__main__":
 import doctest

 doctest.testmod()