code: 
  
  
  
  
  
  
  import json
  import os
  import shutil
  import sys
  import tempfile
  import unittest
  from unittest import TestCase
  from unittest.mock import patch
  
  import faiss
  import numpy as np
  from datasets import Dataset
  
  from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
  from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
  from transformers.integrations import is_ray_available
  from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
  from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
  from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
  from transformers.testing_utils import require_ray
  
  
  sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip
  
  if is_torch_available():
  				from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
  else:
  				SCREAMING_SNAKE_CASE__    :    int								=					None
  
  if is_ray_available():
  				import ray  # noqa: E402 # isort:skip
  				from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
  else:
  				SCREAMING_SNAKE_CASE__    :    Any								=					None
  				SCREAMING_SNAKE_CASE__    :    Tuple								=					None
  				SCREAMING_SNAKE_CASE__    :    List[Any]								=					None
  
  def lowercase						(     SCREAMING_SNAKE_CASE   ) ->       Tuple:
  
  
  			'''simple docstring'''
  
  
  
  
  
  
  
  			if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
  						SCREAMING_SNAKE_CASE_			=	unittest.skip('test requires Datasets, Faiss, psutil'   )(lowerCAmelCase__   )
  			return test_case
  
  
  
  @require_distributed_retrieval
  class 			a_					(    snake_case__    ):
  
  
  
  
  
  	def 	A_(       self			)     ->				int:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	tempfile.mkdtemp()
  				SCREAMING_SNAKE_CASE_			=	8
  
  				# DPR tok
  				SCREAMING_SNAKE_CASE_			=	[
  				    '[UNK]',
  				    '[CLS]',
  				    '[SEP]',
  				    '[PAD]',
  				    '[MASK]',
  				    'want',
  				    '##want',
  				    '##ed',
  				    'wa',
  				    'un',
  				    'runn',
  				    '##ing',
  				    ',',
  				    'low',
  				    'lowest',
  				]
  				SCREAMING_SNAKE_CASE_			=	os.path.join(self.tmpdirname			,    'dpr_tokenizer'			)
  				os.makedirs(SCREAMING_SNAKE_CASE			,    exist_ok=SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	os.path.join(SCREAMING_SNAKE_CASE			,    DPR_VOCAB_FILES_NAMES['vocab_file']			)
  				with open(self.vocab_file			,    'w'			,    encoding='utf-8'			) as vocab_writer:
  							vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]			)			)
  
  				# BART tok
  				SCREAMING_SNAKE_CASE_			=	[
  				    'l',
  				    'o',
  				    'w',
  				    'e',
  				    'r',
  				    's',
  				    't',
  				    'i',
  				    'd',
  				    'n',
  				    '\u0120',
  				    '\u0120l',
  				    '\u0120n',
  				    '\u0120lo',
  				    '\u0120low',
  				    'er',
  				    '\u0120lowest',
  				    '\u0120newer',
  				    '\u0120wider',
  				    '<unk>',
  				]
  				SCREAMING_SNAKE_CASE_			=	dict(zip(SCREAMING_SNAKE_CASE			,    range(len(SCREAMING_SNAKE_CASE			)			)			)			)
  				SCREAMING_SNAKE_CASE_			=	['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
  				SCREAMING_SNAKE_CASE_			=	{'unk_token': '<unk>'}
  
  				SCREAMING_SNAKE_CASE_			=	os.path.join(self.tmpdirname			,    'bart_tokenizer'			)
  				os.makedirs(SCREAMING_SNAKE_CASE			,    exist_ok=SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	os.path.join(SCREAMING_SNAKE_CASE			,    BART_VOCAB_FILES_NAMES['vocab_file']			)
  				SCREAMING_SNAKE_CASE_			=	os.path.join(SCREAMING_SNAKE_CASE			,    BART_VOCAB_FILES_NAMES['merges_file']			)
  				with open(self.vocab_file			,    'w'			,    encoding='utf-8'			) as fp:
  							fp.write(json.dumps(SCREAMING_SNAKE_CASE			) + '\n'			)
  				with open(self.merges_file			,    'w'			,    encoding='utf-8'			) as fp:
  							fp.write('\n'.join(SCREAMING_SNAKE_CASE			)			)
  
  
  
  
  
  	def 	A_(       self			)     ->				str:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname			,    'dpr_tokenizer'			)			)
  
  
  
  
  
  	def 	A_(       self			)     ->				int:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname			,    'bart_tokenizer'			)			)
  
  
  
  
  
  	def 	A_(       self			)     ->				int:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				shutil.rmtree(self.tmpdirname			)
  
  
  
  
  
  	def 	A_(       self			)     ->				Any:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	Dataset.from_dict(
  				    {
  				        'id': ['0', '1'],
  				        'text': ['foo', 'bar'],
  				        'title': ['Foo', 'Bar'],
  				        'embeddings': [np.ones(self.retrieval_vector_size			), 2 * np.ones(self.retrieval_vector_size			)],
  				    }			)
  				dataset.add_faiss_index('embeddings'			,    string_factory='Flat'			,    metric_type=faiss.METRIC_INNER_PRODUCT			)
  				return dataset
  
  
  
  
  
  	def 	A_(       self			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE=1_2345			)     ->				List[Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	self.get_dummy_dataset()
  				SCREAMING_SNAKE_CASE_			=	RagConfig(
  				    retrieval_vector_size=self.retrieval_vector_size			,    question_encoder=DPRConfig().to_dict()			,    generator=BartConfig().to_dict()			,    )
  				with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
  							SCREAMING_SNAKE_CASE_			=	dataset
  							SCREAMING_SNAKE_CASE_			=	RagPyTorchDistributedRetriever(
  							    SCREAMING_SNAKE_CASE			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    )
  							if init_retrieval:
  										retriever.init_retrieval(SCREAMING_SNAKE_CASE			)
  				return retriever
  
  
  
  
  
  	def 	A_(       self			,    SCREAMING_SNAKE_CASE			)     ->				Union[str, Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				ray.init(local_mode=SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	RagConfig(
  				    retrieval_vector_size=self.retrieval_vector_size			,    question_encoder=DPRConfig().to_dict()			,    generator=BartConfig().to_dict()			,    )
  				SCREAMING_SNAKE_CASE_			=	ray.remote(SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	[remote_cls.remote() for _ in range(1			)]
  				with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
  							SCREAMING_SNAKE_CASE_			=	self.get_dummy_dataset()
  							SCREAMING_SNAKE_CASE_			=	RagRayDistributedRetriever(
  							    SCREAMING_SNAKE_CASE			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    retrieval_workers=SCREAMING_SNAKE_CASE			,    )
  							if init_retrieval:
  										retriever.init_retrieval()
  				return retriever
  
  
  
  
  
  	def 	A_(       self			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE=1_2345			)     ->				Dict:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	self.get_dummy_dataset()
  				SCREAMING_SNAKE_CASE_			=	RagConfig(
  				    retrieval_vector_size=self.retrieval_vector_size			,    question_encoder=DPRConfig().to_dict()			,    generator=BartConfig().to_dict()			,    index_name='custom'			,    )
  				if from_disk:
  							SCREAMING_SNAKE_CASE_			=	os.path.join(self.tmpdirname			,    'dataset'			)
  							SCREAMING_SNAKE_CASE_			=	os.path.join(self.tmpdirname			,    'index.faiss'			)
  							dataset.get_index('embeddings'			).save(os.path.join(self.tmpdirname			,    'index.faiss'			)			)
  							dataset.drop_index('embeddings'			)
  							dataset.save_to_disk(os.path.join(self.tmpdirname			,    'dataset'			)			)
  							del dataset
  							SCREAMING_SNAKE_CASE_			=	RagPyTorchDistributedRetriever(
  							    SCREAMING_SNAKE_CASE			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    )
  				else:
  							SCREAMING_SNAKE_CASE_			=	RagPyTorchDistributedRetriever(
  							    SCREAMING_SNAKE_CASE			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    index=CustomHFIndex(config.retrieval_vector_size			,    SCREAMING_SNAKE_CASE			)			,    )
  				if init_retrieval:
  							retriever.init_retrieval(SCREAMING_SNAKE_CASE			)
  				return retriever
  
  
  
  
  
  	def 	A_(       self			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)     ->				Union[str, Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				ray.init(local_mode=SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	self.get_dummy_dataset()
  				SCREAMING_SNAKE_CASE_			=	RagConfig(
  				    retrieval_vector_size=self.retrieval_vector_size			,    question_encoder=DPRConfig().to_dict()			,    generator=BartConfig().to_dict()			,    index_name='custom'			,    )
  				SCREAMING_SNAKE_CASE_			=	ray.remote(SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	[remote_cls.remote() for _ in range(1			)]
  				if from_disk:
  							SCREAMING_SNAKE_CASE_			=	os.path.join(self.tmpdirname			,    'dataset'			)
  							SCREAMING_SNAKE_CASE_			=	os.path.join(self.tmpdirname			,    'index.faiss'			)
  							dataset.get_index('embeddings'			).save(os.path.join(self.tmpdirname			,    'index.faiss'			)			)
  							dataset.drop_index('embeddings'			)
  							dataset.save_to_disk(os.path.join(self.tmpdirname			,    'dataset'			)			)
  							del dataset
  							SCREAMING_SNAKE_CASE_			=	RagRayDistributedRetriever(
  							    SCREAMING_SNAKE_CASE			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    retrieval_workers=SCREAMING_SNAKE_CASE			,    index=CustomHFIndex.load_from_disk(
  							        vector_size=config.retrieval_vector_size			,    dataset_path=config.passages_path			,    index_path=config.index_path			,    )			,    )
  				else:
  							SCREAMING_SNAKE_CASE_			=	RagRayDistributedRetriever(
  							    SCREAMING_SNAKE_CASE			,    question_encoder_tokenizer=self.get_dpr_tokenizer()			,    generator_tokenizer=self.get_bart_tokenizer()			,    retrieval_workers=SCREAMING_SNAKE_CASE			,    index=CustomHFIndex(config.retrieval_vector_size			,    SCREAMING_SNAKE_CASE			)			,    )
  				if init_retrieval:
  							retriever.init_retrieval()
  				return retriever
  
  
  
  
  
  	def 	A_(       self			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)     ->				Optional[Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	retriever.retrieve(SCREAMING_SNAKE_CASE			,    n_docs=SCREAMING_SNAKE_CASE			)
  				self.assertEqual(retrieved_doc_embeds.shape			,    (2, n_docs, self.retrieval_vector_size)			)
  				self.assertEqual(len(SCREAMING_SNAKE_CASE			)			,    2			)
  				self.assertEqual(sorted(doc_dicts[0]			)			,    ['embeddings', 'id', 'text', 'title']			)
  				self.assertEqual(len(doc_dicts[0]['id']			)			,    SCREAMING_SNAKE_CASE			)
  				self.assertEqual(doc_dicts[0]['id'][0]			,    '1'			)  # max inner product is reached with second doc
  				self.assertEqual(doc_dicts[1]['id'][0]			,    '0'			)  # max inner product is reached with first doc
  				self.assertListEqual(doc_ids.tolist()			,    [[1], [0]]			)
  
  
  
  
  
  	def 	A_(       self			)     ->				Any:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	1
  				SCREAMING_SNAKE_CASE_			=	np.array(
  				    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,    dtype=np.floataa			)
  
  				self.distributed_retriever_check(
  				    self.get_dummy_pytorch_distributed_retriever(init_retrieval=SCREAMING_SNAKE_CASE			)			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  
  
  
  
  
  	def 	A_(       self			)     ->				Any:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	1
  				SCREAMING_SNAKE_CASE_			=	np.array(
  				    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,    dtype=np.floataa			)
  
  				self.distributed_retriever_check(
  				    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=SCREAMING_SNAKE_CASE			,    from_disk=SCREAMING_SNAKE_CASE			)			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			,    )
  
  
  
  
  
  	def 	A_(       self			)     ->				Dict:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	1
  				SCREAMING_SNAKE_CASE_			=	np.array(
  				    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,    dtype=np.floataa			)
  
  				self.distributed_retriever_check(
  				    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=SCREAMING_SNAKE_CASE			,    from_disk=SCREAMING_SNAKE_CASE			)			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			,    )
  
  
  
  
  
  	@require_ray
  	def 	A_(       self			)     ->				Optional[Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	1
  				SCREAMING_SNAKE_CASE_			=	np.array(
  				    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,    dtype=np.floataa			)
  
  				self.distributed_retriever_check(
  				    self.get_dummy_ray_distributed_retriever(init_retrieval=SCREAMING_SNAKE_CASE			)			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  				ray.shutdown()
  
  
  
  
  
  	@require_ray
  	def 	A_(       self			)     ->				Any:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	1
  				SCREAMING_SNAKE_CASE_			=	np.array(
  				    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,    dtype=np.floataa			)
  				with self.assertRaises(SCREAMING_SNAKE_CASE			):
  							self.distributed_retriever_check(
  							    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=SCREAMING_SNAKE_CASE			,    from_disk=SCREAMING_SNAKE_CASE			)			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			,    )
  				ray.shutdown()
  
  
  
  
  
  	@require_ray
  	def 	A_(       self			)     ->				Optional[Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	1
  				SCREAMING_SNAKE_CASE_			=	np.array(
  				    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]			,    dtype=np.floataa			)
  
  				self.distributed_retriever_check(
  				    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=SCREAMING_SNAKE_CASE			,    from_disk=SCREAMING_SNAKE_CASE			)			,    SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  				ray.shutdown()
  
style_context: 
  
  
  
  
  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,
  )
  
  
  SCREAMING_SNAKE_CASE__    :    Optional[Any]								=					namedtuple("ExampleNamedTuple", "a b c")
  
  
  
  class 			a_					(    unittest.TestCase    ):
  
  
  
  
  
  	def 	A_(       self			)     ->				List[str]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	torch.randn(5			,    2			)
  				SCREAMING_SNAKE_CASE_			=	torch.device('cuda'			) if torch.cuda.is_available() else torch.device('cpu'			)
  
  				SCREAMING_SNAKE_CASE_			=	send_to_device(SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta.cpu()			,    SCREAMING_SNAKE_CASE			)			)
  
  				SCREAMING_SNAKE_CASE_			=	send_to_device((tensor, [tensor, tensor], 1)			,    SCREAMING_SNAKE_CASE			)
  				self.assertIsInstance(SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta[0].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertIsInstance(resulta[1]			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta[1][0].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertTrue(torch.equal(resulta[1][1].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertEqual(resulta[2]			,    1			)
  
  				SCREAMING_SNAKE_CASE_			=	send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}			,    SCREAMING_SNAKE_CASE			)
  				self.assertIsInstance(SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta['a'].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertIsInstance(resulta['b']			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta['b'][0].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertTrue(torch.equal(resulta['b'][1].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertEqual(resulta['c']			,    1			)
  
  				SCREAMING_SNAKE_CASE_			=	send_to_device(ExampleNamedTuple(a=SCREAMING_SNAKE_CASE			,    b=[tensor, tensor]			,    c=1			)			,    SCREAMING_SNAKE_CASE			)
  				self.assertIsInstance(SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta.a.cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertIsInstance(resulta.b			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta.b[0].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertTrue(torch.equal(resulta.b[1].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertEqual(resulta.c			,    1			)
  
  				SCREAMING_SNAKE_CASE_			=	send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1}			)			,    SCREAMING_SNAKE_CASE			)
  				self.assertIsInstance(SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta['a'].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertIsInstance(resulta['b']			,    SCREAMING_SNAKE_CASE			)
  				self.assertTrue(torch.equal(resulta['b'][0].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertTrue(torch.equal(resulta['b'][1].cpu()			,    SCREAMING_SNAKE_CASE			)			)
  				self.assertEqual(resulta['c']			,    1			)
  
  
  
  
  
  	def 	A_(       self			)     ->				int:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				with self.assertRaises(SCREAMING_SNAKE_CASE			) as cm:
  							SCREAMING_SNAKE_CASE_			=	recursively_apply(torch.tensor			,    (torch.tensor(1			), 1)			,    error_on_other_type=SCREAMING_SNAKE_CASE			)
  				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 	A_(       self			)     ->				int:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	torch.tensor([1, 2, 3, 4, 5]			)
  				self.assertEqual(listify(SCREAMING_SNAKE_CASE			)			,    [1, 2, 3, 4, 5]			)
  
  				SCREAMING_SNAKE_CASE_			=	torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]			)
  				self.assertEqual(listify(SCREAMING_SNAKE_CASE			)			,    [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]			)
  
  				SCREAMING_SNAKE_CASE_			=	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(SCREAMING_SNAKE_CASE			)			,    [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]			)
  
  
  
  
  
  	def 	A_(       self			)     ->				Tuple:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				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 	A_(       self			)     ->				Dict:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	RegressionModel()
  				SCREAMING_SNAKE_CASE_			=	model.forward
  				SCREAMING_SNAKE_CASE_			=	convert_outputs_to_fpaa(model.forward			)
  				SCREAMING_SNAKE_CASE_			=	extract_model_from_parallel(SCREAMING_SNAKE_CASE			,    keep_fpaa_wrapper=SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	pickle.dumps(SCREAMING_SNAKE_CASE			)
  
  
  
  
  
  	@require_cuda
  	def 	A_(       self			)     ->				int:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	RegressionModel()
  				SCREAMING_SNAKE_CASE_			=	model.forward
  				SCREAMING_SNAKE_CASE_			=	torch.cuda.amp.autocast(dtype=torch.floataa			)(model.forward			)
  				SCREAMING_SNAKE_CASE_			=	convert_outputs_to_fpaa(model.forward			)
  				SCREAMING_SNAKE_CASE_			=	extract_model_from_parallel(SCREAMING_SNAKE_CASE			,    keep_fpaa_wrapper=SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	pickle.dumps(SCREAMING_SNAKE_CASE			)
  
  
  
  
  
  	@require_cuda
  	@require_torch_min_version(version='2.0'			)
  	def 	A_(       self			)     ->				Optional[Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	RegressionModel()
  				SCREAMING_SNAKE_CASE_			=	model.forward
  				SCREAMING_SNAKE_CASE_			=	torch.cuda.amp.autocast(dtype=torch.floataa			)(model.forward			)
  				SCREAMING_SNAKE_CASE_			=	convert_outputs_to_fpaa(model.forward			)
  				SCREAMING_SNAKE_CASE_			=	torch.compile(model.forward			,    backend='inductor'			)
  				SCREAMING_SNAKE_CASE_			=	torch.randn(4			,    10			).cuda()
  				SCREAMING_SNAKE_CASE_			=	model(SCREAMING_SNAKE_CASE			)
  
  
  
  
  
  	def 	A_(       self			)     ->				List[str]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	RegressionModel()
  				# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
  				SCREAMING_SNAKE_CASE_			=	torch.nn.parallel.DataParallel(SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	extract_model_from_parallel(SCREAMING_SNAKE_CASE			)
  
  				self.assertEqual(SCREAMING_SNAKE_CASE			,    SCREAMING_SNAKE_CASE			)
  
  
  
  
  
  	@require_torch_min_version(version='2.0'			)
  	def 	A_(       self			)     ->				int:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				SCREAMING_SNAKE_CASE_			=	RegressionModel()
  				SCREAMING_SNAKE_CASE_			=	torch.compile(SCREAMING_SNAKE_CASE			)
  
  				# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
  				SCREAMING_SNAKE_CASE_			=	torch.nn.parallel.DataParallel(SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	torch.compile(SCREAMING_SNAKE_CASE			)
  				SCREAMING_SNAKE_CASE_			=	extract_model_from_parallel(SCREAMING_SNAKE_CASE			)
  
  				self.assertEqual(compiled_model._orig_mod			,    compiled_model_unwrapped._orig_mod			)
  
  
  
  
  
  	def 	A_(       self			)     ->				Optional[Any]:
  
  
  
  
  				"""simple docstring"""
  
  
  
  				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']			)			)
  