





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






   '''simple docstring'''





   def __init__(							self       ):
    super().__init__()
    snake_case__     :					Optional[Any]    =  nn.Linear(3	,		4       )
    snake_case__     :					Union[str, Any]    =  nn.BatchNormad(4       )
    snake_case__     :					Optional[Any]    =  nn.Linear(4	,		5       )


   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE       ):
    return self.lineara(self.batchnorm(self.lineara(__SCREAMING_SNAKE_CASE       )       )       )



class 					__snake_case  (					unittest.TestCase							):






   '''simple docstring'''





   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Tuple    =  ModelForTest()
    with TemporaryDirectory() as tmp_dir:
     offload_state_dict(__SCREAMING_SNAKE_CASE	,		model.state_dict()       )
     snake_case__     :					Tuple    =  os.path.join(__SCREAMING_SNAKE_CASE	,		"""index.json"""       )
     self.assertTrue(os.path.isfile(__SCREAMING_SNAKE_CASE       )       )
     # TODO: add tests on what is inside the index

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

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Dict    =  [torch.floataa, torch.floataa, torch.bfloataa]

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

      snake_case__     :					int    =  load_offloaded_weight(__SCREAMING_SNAKE_CASE	,		index["""weight"""]       )
      self.assertTrue(torch.equal(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Union[str, Any]    =  ModelForTest()
    snake_case__     :					str    =  model.state_dict()
    snake_case__     :					Dict    =  {k: v for k, v in state_dict.items() if """linear2""" not in k}
    snake_case__     :					List[str]    =  {k: v for k, v in state_dict.items() if """linear2""" in k}

    with TemporaryDirectory() as tmp_dir:
     offload_state_dict(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Tuple    =  OffloadedWeightsLoader(state_dict=__SCREAMING_SNAKE_CASE	,		save_folder=__SCREAMING_SNAKE_CASE       )

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

    snake_case__     :					Optional[Any]    =  {k: v for k, v in state_dict.items() if """weight""" in k}
    snake_case__     :					Optional[int]    =  {k: v for k, v in state_dict.items() if """weight""" not in k}

    with TemporaryDirectory() as tmp_dir:
     offload_state_dict(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
     snake_case__     :					Dict    =  OffloadedWeightsLoader(state_dict=__SCREAMING_SNAKE_CASE	,		save_folder=__SCREAMING_SNAKE_CASE       )

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

    with TemporaryDirectory() as tmp_dir:
     offload_state_dict(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
     # Duplicates are removed
     snake_case__     :					List[str]    =  OffloadedWeightsLoader(state_dict=__SCREAMING_SNAKE_CASE	,		save_folder=__SCREAMING_SNAKE_CASE       )

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


   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					str    =  {"""a.1""": 0, """a.10""": 1, """a.2""": 2}
    snake_case__     :					str    =  extract_submodules_state_dict(__SCREAMING_SNAKE_CASE	,		["""a.1""", """a.2"""]       )
    self.assertDictEqual(__SCREAMING_SNAKE_CASE	,		{"""a.1""": 0, """a.2""": 2}       )

    snake_case__     :					Dict    =  {"""a.1.a""": 0, """a.10.a""": 1, """a.2.a""": 2}
    snake_case__     :					Dict    =  extract_submodules_state_dict(__SCREAMING_SNAKE_CASE	,		["""a.1""", """a.2"""]       )
    self.assertDictEqual(__SCREAMING_SNAKE_CASE	,		{"""a.1.a""": 0, """a.2.a""": 2}       )






"""simple docstring"""




import logging

import numpy as np
import pytest
from scipy.linalg import eigh

logging.basicConfig(level=logging.INFO, format='%(message)s')



def 							_lowerCamelCase			(			__a     ):
		return input_array.reshape((input_array.size, 1)     )



def 							_lowerCamelCase			(			__a,				__a,				__a     ):
		SCREAMING_SNAKE_CASE_    			=			np.nan
		for i in range(__a     ):
				SCREAMING_SNAKE_CASE_    			=			features[:, labels == i]
				SCREAMING_SNAKE_CASE_    			=			data.mean(1     )
				# Centralize the data of class i
				SCREAMING_SNAKE_CASE_    			=			data - column_reshape(__a     )
				if i > 0:
						# If covariance_sum is not None
						covariance_sum += np.dot(__a,				centered_data.T     )
				else:
						# If covariance_sum is np.nan (i.e. first loop)
						SCREAMING_SNAKE_CASE_    			=			np.dot(__a,				centered_data.T     )

		return covariance_sum / features.shape[1]



def 							_lowerCamelCase			(			__a,				__a,				__a     ):
		SCREAMING_SNAKE_CASE_    			=			features.mean(1     )
		SCREAMING_SNAKE_CASE_    			=			np.nan
		for i in range(__a     ):
				SCREAMING_SNAKE_CASE_    			=			features[:, labels == i]
				SCREAMING_SNAKE_CASE_    			=			data.shape[1]
				SCREAMING_SNAKE_CASE_    			=			data.mean(1     )
				if i > 0:
						# If covariance_sum is not None
						covariance_sum += device_data * np.dot(
						    column_reshape(__a     ) - column_reshape(__a     ),				(column_reshape(__a     ) - column_reshape(__a     )).T,				)
				else:
						# If covariance_sum is np.nan (i.e. first loop)
						SCREAMING_SNAKE_CASE_    			=			device_data * np.dot(
						    column_reshape(__a     ) - column_reshape(__a     ),				(column_reshape(__a     ) - column_reshape(__a     )).T,				)

		return covariance_sum / features.shape[1]



def 							_lowerCamelCase			(			__a,				__a     ):

		# Check if the features have been loaded
		if features.any():
				SCREAMING_SNAKE_CASE_    			=			features.mean(1     )
				# Center the dataset
				SCREAMING_SNAKE_CASE_    			=			features - np.reshape(__a,				(data_mean.size, 1)     )
				SCREAMING_SNAKE_CASE_    			=			np.dot(__a,				centered_data.T     ) / features.shape[1]
				SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			np.linalg.eigh(__a     )
				# Take all the columns in the reverse order (-1), and then takes only the first
				SCREAMING_SNAKE_CASE_    			=			eigenvectors[:, ::-1][:, 0:dimensions]
				# Project the database on the new space
				SCREAMING_SNAKE_CASE_    			=			np.dot(filtered_eigenvectors.T,				__a     )
				logging.info('''Principal Component Analysis computed'''     )

				return projected_data
		else:
				logging.basicConfig(level=logging.ERROR,				format='''%(message)s''',				force=__a     )
				logging.error('''Dataset empty'''     )
				raise AssertionError



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a     ):
		assert classes > dimensions

		# Check if features have been already loaded
		if features.any:
				SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			eigh(
				    covariance_between_classes(__a,				__a,				__a     ),				covariance_within_classes(__a,				__a,				__a     ),				)
				SCREAMING_SNAKE_CASE_    			=			eigenvectors[:, ::-1][:, :dimensions]
				SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			np.linalg.svd(__a     )
				SCREAMING_SNAKE_CASE_    			=			svd_matrix[:, 0:dimensions]
				SCREAMING_SNAKE_CASE_    			=			np.dot(filtered_svd_matrix.T,				__a     )
				logging.info('''Linear Discriminant Analysis computed'''     )

				return projected_data
		else:
				logging.basicConfig(level=logging.ERROR,				format='''%(message)s''',				force=__a     )
				logging.error('''Dataset empty'''     )
				raise AssertionError



def 							_lowerCamelCase			(			):
		# Create dummy dataset with 2 classes and 3 features
		SCREAMING_SNAKE_CASE_    			=			np.array([[1, 2, 3, 4, 5], [2, 3, 4, 5, 6], [3, 4, 5, 6, 7]]     )
		SCREAMING_SNAKE_CASE_    			=			np.array([0, 0, 0, 1, 1]     )
		SCREAMING_SNAKE_CASE_    			=			2
		SCREAMING_SNAKE_CASE_    			=			2

		# Assert that the function raises an AssertionError if dimensions > classes
		with pytest.raises(__a     ) as error_info:
				SCREAMING_SNAKE_CASE_    			=			linear_discriminant_analysis(
				    __a,				__a,				__a,				__a     )
				if isinstance(__a,				np.ndarray     ):
						raise AssertionError(
						    '''Did not raise AssertionError for dimensions > classes'''     )
				assert error_info.type is AssertionError



def 							_lowerCamelCase			(			):
		SCREAMING_SNAKE_CASE_    			=			np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]     )
		SCREAMING_SNAKE_CASE_    			=			2
		SCREAMING_SNAKE_CASE_    			=			np.array([[6.9_2_8_2_0_3_2_3, 8.6_6_0_2_5_4_0_4, 1_0.3_9_2_3_0_4_8_5], [3.0, 3.0, 3.0]]     )

		with pytest.raises(__a     ) as error_info:
				SCREAMING_SNAKE_CASE_    			=			principal_component_analysis(__a,				__a     )
				if not np.allclose(__a,				__a     ):
						raise AssertionError
				assert error_info.type is AssertionError


if __name__ == "__main__":
						import doctest

						doctest.testmod()