
"""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,
)


__A					=			namedtuple('ExampleNamedTuple', 'a b c')







class 						lowerCamelCase__   (   unittest.TestCase							):





			'''simple docstring'''



			def    _lowerCAmelCase						(				self    )					->	Tuple:
				_lowerCAmelCase								=torch.randn(5	,	2    )
				_lowerCAmelCase								=torch.device("""cuda"""    ) if torch.cuda.is_available() else torch.device("""cpu"""    )

				_lowerCAmelCase								=send_to_device(__UpperCAmelCase	,	__UpperCAmelCase    )
				self.assertTrue(torch.equal(resulta.cpu()	,	__UpperCAmelCase    )    )

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

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

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

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



			def    _lowerCAmelCase						(				self    )					->	List[str]:
				with self.assertRaises(__UpperCAmelCase    ) as cm:
					_lowerCAmelCase								=recursively_apply(torch.tensor	,	(torch.tensor(1    ), 1)	,	error_on_other_type=__UpperCAmelCase    )
				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    _lowerCAmelCase						(				self    )					->	Optional[Any]:
				_lowerCAmelCase								=torch.tensor([1, 2, 3, 4, 5]    )
				self.assertEqual(listify(__UpperCAmelCase    )	,	[1, 2, 3, 4, 5]    )

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

				_lowerCAmelCase								=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(__UpperCAmelCase    )	,	[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]    )



			def    _lowerCAmelCase						(				self    )					->	Optional[Any]:
				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    _lowerCAmelCase						(				self    )					->	List[str]:
				_lowerCAmelCase								=RegressionModel()
				_lowerCAmelCase								=model.forward
				_lowerCAmelCase								=convert_outputs_to_fpaa(model.forward    )
				_lowerCAmelCase								=extract_model_from_parallel(__UpperCAmelCase	,	keep_fpaa_wrapper=__UpperCAmelCase    )
				_lowerCAmelCase								=pickle.dumps(__UpperCAmelCase    )



			@require_cuda
			def    _lowerCAmelCase						(				self    )					->	Dict:
				_lowerCAmelCase								=RegressionModel()
				_lowerCAmelCase								=model.forward
				_lowerCAmelCase								=torch.cuda.amp.autocast(dtype=torch.floataa    )(model.forward    )
				_lowerCAmelCase								=convert_outputs_to_fpaa(model.forward    )
				_lowerCAmelCase								=extract_model_from_parallel(__UpperCAmelCase	,	keep_fpaa_wrapper=__UpperCAmelCase    )
				_lowerCAmelCase								=pickle.dumps(__UpperCAmelCase    )



			@require_cuda
			@require_torch_min_version(version="""2.0"""    )
			def    _lowerCAmelCase						(				self    )					->	str:
				_lowerCAmelCase								=RegressionModel()
				_lowerCAmelCase								=model.forward
				_lowerCAmelCase								=torch.cuda.amp.autocast(dtype=torch.floataa    )(model.forward    )
				_lowerCAmelCase								=convert_outputs_to_fpaa(model.forward    )
				_lowerCAmelCase								=torch.compile(model.forward	,	backend="""inductor"""    )
				_lowerCAmelCase								=torch.randn(4	,	10    ).cuda()
				_lowerCAmelCase								=model(__UpperCAmelCase    )



			def    _lowerCAmelCase						(				self    )					->	Tuple:
				_lowerCAmelCase								=RegressionModel()
				# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
				_lowerCAmelCase								=torch.nn.parallel.DataParallel(__UpperCAmelCase    )
				_lowerCAmelCase								=extract_model_from_parallel(__UpperCAmelCase    )

				self.assertEqual(__UpperCAmelCase	,	__UpperCAmelCase    )



			@require_torch_min_version(version="""2.0"""    )
			def    _lowerCAmelCase						(				self    )					->	List[str]:
				_lowerCAmelCase								=RegressionModel()
				_lowerCAmelCase								=torch.compile(__UpperCAmelCase    )

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

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



			def    _lowerCAmelCase						(				self    )					->	Tuple:
				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"""


from ...configuration_utils import PretrainedConfig
from ...utils import logging


__A					=			logging.get_logger(__name__)

__A					=			{
    'uw-madison/nystromformer-512': 'https://huggingface.co/uw-madison/nystromformer-512/resolve/main/config.json',
    # See all Nystromformer models at https://huggingface.co/models?filter=nystromformer
}







class 						lowerCamelCase__   (   __magic_name__							):





			'''simple docstring'''
			lowerCamelCase       =						'''nystromformer'''



			def __init__(				self	,	__UpperCAmelCase=3_00_00	,	__UpperCAmelCase=7_68	,	__UpperCAmelCase=12	,	__UpperCAmelCase=12	,	__UpperCAmelCase=30_72	,	__UpperCAmelCase="gelu_new"	,	__UpperCAmelCase=0.1	,	__UpperCAmelCase=0.1	,	__UpperCAmelCase=5_10	,	__UpperCAmelCase=2	,	__UpperCAmelCase=64	,	__UpperCAmelCase=64	,	__UpperCAmelCase=65	,	__UpperCAmelCase=False	,	__UpperCAmelCase=0.0_2	,	__UpperCAmelCase=1e-5	,	__UpperCAmelCase=1	,	__UpperCAmelCase=0	,	__UpperCAmelCase=2	,	**__UpperCAmelCase	,	)					->	Union[str, Any]:
				_lowerCAmelCase								=vocab_size
				_lowerCAmelCase								=max_position_embeddings
				_lowerCAmelCase								=hidden_size
				_lowerCAmelCase								=num_hidden_layers
				_lowerCAmelCase								=num_attention_heads
				_lowerCAmelCase								=intermediate_size
				_lowerCAmelCase								=hidden_act
				_lowerCAmelCase								=hidden_dropout_prob
				_lowerCAmelCase								=attention_probs_dropout_prob
				_lowerCAmelCase								=initializer_range
				_lowerCAmelCase								=type_vocab_size
				_lowerCAmelCase								=segment_means_seq_len
				_lowerCAmelCase								=num_landmarks
				_lowerCAmelCase								=conv_kernel_size
				_lowerCAmelCase								=inv_coeff_init_option
				_lowerCAmelCase								=layer_norm_eps
				super().__init__(pad_token_id=__UpperCAmelCase	,	bos_token_id=__UpperCAmelCase	,	eos_token_id=__UpperCAmelCase	,	**__UpperCAmelCase    )



