





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





         """simple docstring"""



         super().__init__()
         __snake_case							       =   nn.Linear(3				,			4      )
         __snake_case							       =   nn.BatchNormad(4      )
         __snake_case							       =   nn.Linear(4				,			5      )


      def    a (self       :      Optional[int]				,			a__       :      Optional[Any]      ):





         """simple docstring"""



         return self.lineara(self.batchnorm(self.lineara(lowerCAmelCase__      )      )      )






class 		SCREAMING_SNAKE_CASE__				(       unittest.TestCase  ):


      def    a (self       :      List[Any]      ):





         """simple docstring"""



         __snake_case							       =   ModelForTest()
         with TemporaryDirectory() as tmp_dir:
            offload_state_dict(lowerCAmelCase__				,			model.state_dict()      )
            __snake_case							       =   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"]:
               __snake_case							       =   os.path.join(lowerCAmelCase__				,			f"""{key}.dat"""      )
               self.assertTrue(os.path.isfile(lowerCAmelCase__      )      )
               # TODO: add tests on the fact weights are properly loaded


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   [torch.floataa, torch.floataa, torch.bfloataa]

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

               __snake_case							       =   load_offloaded_weight(lowerCAmelCase__				,			index['''weight''']      )
               self.assertTrue(torch.equal(lowerCAmelCase__				,			lowerCAmelCase__      )      )


      def    a (self       :      str      ):





         """simple docstring"""



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

         with TemporaryDirectory() as tmp_dir:
            offload_state_dict(lowerCAmelCase__				,			lowerCAmelCase__      )
            __snake_case							       =   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]      )      )

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

         with TemporaryDirectory() as tmp_dir:
            offload_state_dict(lowerCAmelCase__				,			lowerCAmelCase__      )
            __snake_case							       =   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
            __snake_case							       =   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    a (self       :      List[Any]      ):





         """simple docstring"""



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

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



"""simple docstring"""





import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs



@pytest.mark.parametrize(
    """kwargs, expected"""  ,				[
        ({"""num_shards""": 0, """max_num_jobs""": 1}, []),
        ({"""num_shards""": 10, """max_num_jobs""": 1}, [range(10				)]),
        ({"""num_shards""": 10, """max_num_jobs""": 10}, [range(UpperCamelCase_  ,				i + 1				) for i in range(10				)]),
        ({"""num_shards""": 1, """max_num_jobs""": 10}, [range(1				)]),
        ({"""num_shards""": 10, """max_num_jobs""": 3}, [range(0  ,				4				), range(4  ,				7				), range(7  ,				10				)]),
        ({"""num_shards""": 3, """max_num_jobs""": 10}, [range(0  ,				1				), range(1  ,				2				), range(2  ,				3				)]),
    ]  ,				)
def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_				):
		__SCREAMING_SNAKE_CASE            = _distribute_shards(**UpperCamelCase_				)
		assert out == expected



@pytest.mark.parametrize(
    """gen_kwargs, max_num_jobs, expected"""  ,				[
        ({"""foo""": 0}, 10, [{"""foo""": 0}]),
        ({"""shards""": [0, 1, 2, 3]}, 1, [{"""shards""": [0, 1, 2, 3]}]),
        ({"""shards""": [0, 1, 2, 3]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}, {"""shards""": [2]}, {"""shards""": [3]}]),
        ({"""shards""": [0, 1]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}]),
        ({"""shards""": [0, 1, 2, 3]}, 2, [{"""shards""": [0, 1]}, {"""shards""": [2, 3]}]),
    ]  ,				)
def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_  ,				UpperCamelCase_				):
		__SCREAMING_SNAKE_CASE            = _split_gen_kwargs(UpperCamelCase_  ,				UpperCamelCase_				)
		assert out == expected



@pytest.mark.parametrize(
    """gen_kwargs, expected"""  ,				[
        ({"""foo""": 0}, 1),
        ({"""shards""": [0]}, 1),
        ({"""shards""": [0, 1, 2, 3]}, 4),
        ({"""shards""": [0, 1, 2, 3], """foo""": 0}, 4),
        ({"""shards""": [0, 1, 2, 3], """other""": (0, 1)}, 4),
        ({"""shards""": [0, 1, 2, 3], """shards2""": [0, 1]}, RuntimeError),
    ]  ,				)
def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_				):
		if expected is RuntimeError:
				with pytest.raises(UpperCamelCase_				):
						_number_of_shards_in_gen_kwargs(UpperCamelCase_				)
		else:
				__SCREAMING_SNAKE_CASE            = _number_of_shards_in_gen_kwargs(UpperCamelCase_				)
				assert out == expected




