






'''simple docstring'''




import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    FlaubertForMultipleChoice,
		    FlaubertForQuestionAnswering,
		    FlaubertForQuestionAnsweringSimple,
		    FlaubertForSequenceClassification,
		    FlaubertForTokenClassification,
		    FlaubertModel,
		    FlaubertWithLMHeadModel,
		)
		from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class       __lowercase					(__lowerCamelCase    ):

	def __init__(      self   :						Optional[Any]  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						Optional[int]=13  ,					UpperCAmelCase_   :						Optional[int]=7  ,					UpperCAmelCase_   :						Any=True  ,					UpperCAmelCase_   :						List[Any]=True  ,					UpperCAmelCase_   :						Tuple=True  ,					UpperCAmelCase_   :						Union[str, Any]=True  ,					UpperCAmelCase_   :						str=True  ,					UpperCAmelCase_   :						Optional[Any]=False  ,					UpperCAmelCase_   :						List[Any]=False  ,					UpperCAmelCase_   :						Optional[Any]=False  ,					UpperCAmelCase_   :						Optional[int]=2  ,					UpperCAmelCase_   :						int=99  ,					UpperCAmelCase_   :						List[str]=0  ,					UpperCAmelCase_   :						List[str]=32  ,					UpperCAmelCase_   :						int=5  ,					UpperCAmelCase_   :						Any=4  ,					UpperCAmelCase_   :						Dict=0.1  ,					UpperCAmelCase_   :						int=0.1  ,					UpperCAmelCase_   :						List[str]=512  ,					UpperCAmelCase_   :						Tuple=12  ,					UpperCAmelCase_   :						str=2  ,					UpperCAmelCase_   :						Tuple=0.02  ,					UpperCAmelCase_   :						Any=3  ,					UpperCAmelCase_   :						Optional[int]=4  ,					UpperCAmelCase_   :						Union[str, Any]="last"  ,					UpperCAmelCase_   :						Any=None  ,					UpperCAmelCase_   :						int=None  ,					):
			UpperCamelCase__				:					int					=					parent
			UpperCamelCase__				:					Any					=					batch_size
			UpperCamelCase__				:					Tuple					=					seq_length
			UpperCamelCase__				:					Tuple					=					is_training
			UpperCamelCase__				:					Union[str, Any]					=					use_input_lengths
			UpperCamelCase__				:					Any					=					use_token_type_ids
			UpperCamelCase__				:					List[Any]					=					use_labels
			UpperCamelCase__				:					Optional[Any]					=					gelu_activation
			UpperCamelCase__				:					Union[str, Any]					=					sinusoidal_embeddings
			UpperCamelCase__				:					Union[str, Any]					=					causal
			UpperCamelCase__				:					Optional[Any]					=					asm
			UpperCamelCase__				:					Union[str, Any]					=					n_langs
			UpperCamelCase__				:					str					=					vocab_size
			UpperCamelCase__				:					int					=					n_special
			UpperCamelCase__				:					Dict					=					hidden_size
			UpperCamelCase__				:					List[str]					=					num_hidden_layers
			UpperCamelCase__				:					Tuple					=					num_attention_heads
			UpperCamelCase__				:					str					=					hidden_dropout_prob
			UpperCamelCase__				:					Optional[int]					=					attention_probs_dropout_prob
			UpperCamelCase__				:					int					=					max_position_embeddings
			UpperCamelCase__				:					Optional[Any]					=					type_vocab_size
			UpperCamelCase__				:					Optional[Any]					=					type_sequence_label_size
			UpperCamelCase__				:					Dict					=					initializer_range
			UpperCamelCase__				:					int					=					num_labels
			UpperCamelCase__				:					Optional[Any]					=					num_choices
			UpperCamelCase__				:					Dict					=					summary_type
			UpperCamelCase__				:					int					=					use_proj
			UpperCamelCase__				:					Optional[int]					=					scope

	def      __UpperCamelCase							(      self   :						Dict):
			UpperCamelCase__				:					int					=					ids_tensor([self.batch_size, self.seq_length]  ,					self.vocab_size)
			UpperCamelCase__				:					List[str]					=					random_attention_mask([self.batch_size, self.seq_length])

			UpperCamelCase__				:					Any					=					None
			if self.use_input_lengths:
					UpperCamelCase__				:					Union[str, Any]					=					(
					    ids_tensor([self.batch_size]  ,					vocab_size=2) + self.seq_length - 2
					)  # small variation of seq_length

			UpperCamelCase__				:					Optional[int]					=					None
			if self.use_token_type_ids:
					UpperCamelCase__				:					Optional[int]					=					ids_tensor([self.batch_size, self.seq_length]  ,					self.n_langs)

			UpperCamelCase__				:					str					=					None
			UpperCamelCase__				:					List[str]					=					None
			UpperCamelCase__				:					int					=					None
			if self.use_labels:
					UpperCamelCase__				:					str					=					ids_tensor([self.batch_size]  ,					self.type_sequence_label_size)
					UpperCamelCase__				:					List[str]					=					ids_tensor([self.batch_size, self.seq_length]  ,					self.num_labels)
					UpperCamelCase__				:					List[str]					=					ids_tensor([self.batch_size]  ,					2).float()
					UpperCamelCase__				:					str					=					ids_tensor([self.batch_size]  ,					self.num_choices)

			UpperCamelCase__				:					Tuple					=					self.get_config()

			return (
			    config,
			    input_ids,
			    token_type_ids,
			    input_lengths,
			    sequence_labels,
			    token_labels,
			    is_impossible_labels,
			    choice_labels,
			    input_mask,
			)

	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			return FlaubertConfig(
			    vocab_size=self.vocab_size  ,					n_special=self.n_special  ,					emb_dim=self.hidden_size  ,					n_layers=self.num_hidden_layers  ,					n_heads=self.num_attention_heads  ,					dropout=self.hidden_dropout_prob  ,					attention_dropout=self.attention_probs_dropout_prob  ,					gelu_activation=self.gelu_activation  ,					sinusoidal_embeddings=self.sinusoidal_embeddings  ,					asm=self.asm  ,					causal=self.causal  ,					n_langs=self.n_langs  ,					max_position_embeddings=self.max_position_embeddings  ,					initializer_range=self.initializer_range  ,					summary_type=self.summary_type  ,					use_proj=self.use_proj  ,					)

	def      __UpperCamelCase							(      self   :						Any  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Any  ,					):
			UpperCamelCase__				:					str					=					FlaubertModel(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					Union[str, Any]					=					model(UpperCAmelCase_  ,					lengths=UpperCAmelCase_  ,					langs=UpperCAmelCase_)
			UpperCamelCase__				:					Optional[Any]					=					model(UpperCAmelCase_  ,					langs=UpperCAmelCase_)
			UpperCamelCase__				:					Any					=					model(UpperCAmelCase_)
			self.parent.assertEqual(result.last_hidden_state.shape  ,					(self.batch_size, self.seq_length, self.hidden_size))

	def      __UpperCamelCase							(      self   :						str  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						Dict  ,					):
			UpperCamelCase__				:					List[Any]					=					FlaubertWithLMHeadModel(UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()

			UpperCamelCase__				:					Optional[int]					=					model(UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					labels=UpperCAmelCase_)
			self.parent.assertEqual(result.loss.shape  ,					())
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, self.seq_length, self.vocab_size))

	def      __UpperCamelCase							(      self   :						List[str]  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						Optional[int]  ,					):
			UpperCamelCase__				:					List[Any]					=					FlaubertForQuestionAnsweringSimple(UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()

			UpperCamelCase__				:					List[Any]					=					model(UpperCAmelCase_)

			UpperCamelCase__				:					List[str]					=					model(UpperCAmelCase_  ,					start_positions=UpperCAmelCase_  ,					end_positions=UpperCAmelCase_)
			self.parent.assertEqual(result.start_logits.shape  ,					(self.batch_size, self.seq_length))
			self.parent.assertEqual(result.end_logits.shape  ,					(self.batch_size, self.seq_length))

	def      __UpperCamelCase							(      self   :						int  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						List[str]  ,					):
			UpperCamelCase__				:					str					=					FlaubertForQuestionAnswering(UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()

			UpperCamelCase__				:					Optional[int]					=					model(UpperCAmelCase_)

			UpperCamelCase__				:					Optional[Any]					=					model(
			    UpperCAmelCase_  ,					start_positions=UpperCAmelCase_  ,					end_positions=UpperCAmelCase_  ,					cls_index=UpperCAmelCase_  ,					is_impossible=UpperCAmelCase_  ,					p_mask=UpperCAmelCase_  ,					)

			UpperCamelCase__				:					Optional[int]					=					model(
			    UpperCAmelCase_  ,					start_positions=UpperCAmelCase_  ,					end_positions=UpperCAmelCase_  ,					cls_index=UpperCAmelCase_  ,					is_impossible=UpperCAmelCase_  ,					)

			((UpperCamelCase__),	)				:					Tuple					=					result_with_labels.to_tuple()

			UpperCamelCase__				:					int					=					model(UpperCAmelCase_  ,					start_positions=UpperCAmelCase_  ,					end_positions=UpperCAmelCase_)

			((UpperCamelCase__),	)				:					int					=					result_with_labels.to_tuple()

			self.parent.assertEqual(result_with_labels.loss.shape  ,					())
			self.parent.assertEqual(result.start_top_log_probs.shape  ,					(self.batch_size, model.config.start_n_top))
			self.parent.assertEqual(result.start_top_index.shape  ,					(self.batch_size, model.config.start_n_top))
			self.parent.assertEqual(
			    result.end_top_log_probs.shape  ,					(self.batch_size, model.config.start_n_top * model.config.end_n_top))
			self.parent.assertEqual(
			    result.end_top_index.shape  ,					(self.batch_size, model.config.start_n_top * model.config.end_n_top))
			self.parent.assertEqual(result.cls_logits.shape  ,					(self.batch_size,))

	def      __UpperCamelCase							(      self   :						Optional[int]  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						int  ,					):
			UpperCamelCase__				:					str					=					FlaubertForSequenceClassification(UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()

			UpperCamelCase__				:					Tuple					=					model(UpperCAmelCase_)
			UpperCamelCase__				:					Any					=					model(UpperCAmelCase_  ,					labels=UpperCAmelCase_)

			self.parent.assertEqual(result.loss.shape  ,					())
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, self.type_sequence_label_size))

	def      __UpperCamelCase							(      self   :						Optional[int]  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						Optional[Any]  ,					):
			UpperCamelCase__				:					List[Any]					=					self.num_labels
			UpperCamelCase__				:					Optional[Any]					=					FlaubertForTokenClassification(UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()

			UpperCamelCase__				:					Union[str, Any]					=					model(UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					labels=UpperCAmelCase_)
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, self.seq_length, self.num_labels))

	def      __UpperCamelCase							(      self   :						Tuple  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						int  ,					):
			UpperCamelCase__				:					Optional[Any]					=					self.num_choices
			UpperCamelCase__				:					List[Any]					=					FlaubertForMultipleChoice(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					List[str]					=					input_ids.unsqueeze(1).expand(-1  ,					self.num_choices  ,					-1).contiguous()
			UpperCamelCase__				:					List[str]					=					token_type_ids.unsqueeze(1).expand(-1  ,					self.num_choices  ,					-1).contiguous()
			UpperCamelCase__				:					Tuple					=					input_mask.unsqueeze(1).expand(-1  ,					self.num_choices  ,					-1).contiguous()
			UpperCamelCase__				:					List[str]					=					model(
			    UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					labels=UpperCAmelCase_  ,					)
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, self.num_choices))





	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			UpperCamelCase__				:					List[Any]					=					self.prepare_config_and_inputs()
			(
			    (
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	
			)				:					str					=					config_and_inputs
			UpperCamelCase__				:					Optional[Any]					=					{
			    'input_ids': input_ids,
			    'token_type_ids': token_type_ids,
			    'lengths': input_lengths,
			    'attention_mask': input_mask,
			}
			return config, inputs_dict







@require_torch
class       __lowercase					(__lowerCamelCase   ,				__lowerCamelCase   ,				unittest.TestCase    ):
	_lowerCamelCase              =  (
	    (
	        FlaubertModel,
	        FlaubertWithLMHeadModel,
	        FlaubertForQuestionAnswering,
	        FlaubertForQuestionAnsweringSimple,
	        FlaubertForSequenceClassification,
	        FlaubertForTokenClassification,
	        FlaubertForMultipleChoice,
	    )
	    if is_torch_available()
	    else ()
	)
	_lowerCamelCase              =  (
	    {
	        '''feature-extraction''': FlaubertModel,
	        '''fill-mask''': FlaubertWithLMHeadModel,
	        '''question-answering''': FlaubertForQuestionAnsweringSimple,
	        '''text-classification''': FlaubertForSequenceClassification,
	        '''token-classification''': FlaubertForTokenClassification,
	        '''zero-shot''': FlaubertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)

	def      __UpperCamelCase							(      self   :						Tuple  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						Union[str, Any]):
			if (
			    pipeline_test_casse_name == "QAPipelineTests"
			    and tokenizer_name is not None
			    and not tokenizer_name.endswith('Fast')
			):
					# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
					# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
					# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
					return True

			return False

	def      __UpperCamelCase							(      self   :						Optional[Any]  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Union[str, Any]=False):
			UpperCamelCase__				:					str					=					super()._prepare_for_class(UpperCAmelCase_  ,					UpperCAmelCase_  ,					return_labels=UpperCAmelCase_)

			if return_labels:
					if model_class.__name__ == "FlaubertForQuestionAnswering":
							UpperCamelCase__				:					int					=					torch.zeros(
							    self.model_tester.batch_size  ,					dtype=torch.long  ,					device=UpperCAmelCase_)
							UpperCamelCase__				:					Union[str, Any]					=					torch.zeros(
							    self.model_tester.batch_size  ,					dtype=torch.long  ,					device=UpperCAmelCase_)

			return inputs_dict

	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			UpperCamelCase__				:					str					=					FlaubertModelTester(self)
			UpperCamelCase__				:					int					=					ConfigTester(self  ,					config_class=UpperCAmelCase_  ,					emb_dim=37)

	def      __UpperCamelCase							(      self   :						Optional[Any]):
			self.config_tester.run_common_tests()

	def      __UpperCamelCase							(      self   :						str):
			UpperCamelCase__				:					str					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_flaubert_model(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Any):
			UpperCamelCase__				:					Union[str, Any]					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_flaubert_lm_head(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Optional[int]):
			UpperCamelCase__				:					Dict					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_flaubert_simple_qa(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Optional[int]):
			UpperCamelCase__				:					int					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_flaubert_qa(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Optional[Any]):
			UpperCamelCase__				:					Tuple					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_flaubert_sequence_classif(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			UpperCamelCase__				:					Dict					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_flaubert_token_classif(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Optional[int]):
			UpperCamelCase__				:					Dict					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_flaubert_multiple_choice(*UpperCAmelCase_)

	@slow
	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					UpperCamelCase__				:					int					=					FlaubertModel.from_pretrained(UpperCAmelCase_)
					self.assertIsNotNone(UpperCAmelCase_)





	@slow
	@require_torch_gpu
	def      __UpperCamelCase							(      self   :						Optional[int]):
			UpperCamelCase__,	UpperCamelCase__				:					List[Any]					=					self.model_tester.prepare_config_and_inputs_for_common()
			for model_class in self.all_model_classes:
					# FlauBertForMultipleChoice behaves incorrectly in JIT environments.
					if model_class == FlaubertForMultipleChoice:
							return

					UpperCamelCase__				:					str					=					True
					UpperCamelCase__				:					List[Any]					=					model_class(config=UpperCAmelCase_)

					UpperCamelCase__				:					Tuple					=					self._prepare_for_class(UpperCAmelCase_  ,					UpperCAmelCase_)
					UpperCamelCase__				:					Optional[int]					=					torch.jit.trace(
					    UpperCAmelCase_  ,					(inputs_dict['input_ids'].to('cpu'), inputs_dict['attention_mask'].to('cpu')))

					with tempfile.TemporaryDirectory() as tmp:
							torch.jit.save(UpperCAmelCase_  ,					os.path.join(UpperCAmelCase_  ,					'traced_model.pt'))
							UpperCamelCase__				:					List[str]					=					torch.jit.load(os.path.join(UpperCAmelCase_  ,					'traced_model.pt')  ,					map_location=UpperCAmelCase_)
							loaded(inputs_dict['input_ids'].to(UpperCAmelCase_)  ,					inputs_dict['attention_mask'].to(UpperCAmelCase_))







@require_torch
class       __lowercase					(unittest.TestCase    ):

	@slow
	def      __UpperCamelCase							(      self   :						List[Any]):
			UpperCamelCase__				:					str					=					FlaubertModel.from_pretrained('flaubert/flaubert_base_cased')
			UpperCamelCase__				:					Dict					=					torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]])
			with torch.no_grad():
					UpperCamelCase__				:					str					=					model(UpperCAmelCase_)[0]
			UpperCamelCase__				:					Optional[int]					=					torch.Size((1, 11, 768))
			self.assertEqual(output.shape  ,					UpperCAmelCase_)
			UpperCamelCase__				:					Optional[Any]					=					torch.tensor(
			    [[[-2.62_51, -1.42_98, -0.02_27], [-2.85_10, -1.63_87, 0.22_58], [-2.81_14, -1.18_32, -0.30_66]]])

			self.assertTrue(torch.allclose(output[:, :3, :3]  ,					UpperCAmelCase_  ,					atol=1e-4))












'''simple docstring'''




import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map







def       __UpperCAmelCase						(  lowerCamelCase_)							->     List[Tuple[int, ...]]:
		UpperCamelCase__				:					int					=					[]
		if isinstance(lowerCamelCase_     ,      lowerCamelCase_):
				for v in tree.values():
						shapes.extend(_fetch_dims(lowerCamelCase_))
		elif isinstance(lowerCamelCase_     ,      (list, tuple)):
				for t in tree:
						shapes.extend(_fetch_dims(lowerCamelCase_))
		elif isinstance(lowerCamelCase_     ,      torch.Tensor):
				shapes.append(tree.shape)
		else:
				raise ValueError('Not supported')

		return shapes







@torch.jit.ignore
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_)							->     Tuple[int, ...]:
		UpperCamelCase__				:					int					=					[]
		for d in reversed(lowerCamelCase_):
				idx.append(flat_idx % d)
				UpperCamelCase__				:					Any					=					flat_idx // d

		return tuple(reversed(lowerCamelCase_))







@torch.jit.ignore
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_ = None     ,      lowerCamelCase_ = None     ,      )							->     List[Tuple[slice, ...]]:

		# start_edges and end_edges both indicate whether, starting from any given
		# dimension, the start/end index is at the top/bottom edge of the
		# corresponding tensor, modeled as a tree
		def reduce_edge_list(lowerCamelCase_) -> None:
				UpperCamelCase__				:					Tuple					=					True
				for i in range(len(lowerCamelCase_)):
						UpperCamelCase__				:					List[Any]					=					-1 * (i + 1)
						l[reversed_idx] &= tally
						UpperCamelCase__				:					Optional[Any]					=					l[reversed_idx]

		if start_edges is None:
				UpperCamelCase__				:					int					=					[s == 0 for s in start]
				reduce_edge_list(lowerCamelCase_)
		if end_edges is None:
				UpperCamelCase__				:					List[str]					=					[e == (d - 1) for e, d in zip(lowerCamelCase_     ,      lowerCamelCase_)]
				reduce_edge_list(lowerCamelCase_)

		# Base cases. Either start/end are empty and we're done, or the final,
		# one-dimensional tensor can be simply sliced
		if len(lowerCamelCase_) == 0:
				return [()]
		elif len(lowerCamelCase_) == 1:
				return [(slice(start[0]     ,      end[0] + 1),)]

		UpperCamelCase__				:					List[Tuple[slice, ...]]					=					[]
		UpperCamelCase__				:					List[slice]					=					[]

		# Dimensions common to start and end can be selected directly
		for s, e in zip(lowerCamelCase_     ,      lowerCamelCase_):
				if s == e:
						path_list.append(slice(lowerCamelCase_     ,      s + 1))
				else:
						break

		UpperCamelCase__				:					Tuple[slice, ...]					=					tuple(lowerCamelCase_)
		UpperCamelCase__				:					Dict					=					len(lowerCamelCase_)

		# start == end, and we're done
		if divergence_idx == len(lowerCamelCase_):
				return [path]

		def upper() -> Tuple[Tuple[slice, ...], ...]:
				assert start_edges is not None
				assert end_edges is not None

				UpperCamelCase__				:					str					=					start[divergence_idx]
				return tuple(
				    path + (slice(lowerCamelCase_     ,      sdi + 1),) + s
				    for s in _get_minimal_slice_set(
				        start[divergence_idx + 1 :]     ,      [d - 1 for d in dims[divergence_idx + 1 :]]     ,      dims[divergence_idx + 1 :]     ,      start_edges=start_edges[divergence_idx + 1 :]     ,      end_edges=[True for _ in end_edges[divergence_idx + 1 :]]     ,      ))

		def lower() -> Tuple[Tuple[slice, ...], ...]:
				assert start_edges is not None
				assert end_edges is not None

				UpperCamelCase__				:					Optional[int]					=					end[divergence_idx]
				return tuple(
				    path + (slice(lowerCamelCase_     ,      edi + 1),) + s
				    for s in _get_minimal_slice_set(
				        [0 for _ in start[divergence_idx + 1 :]]     ,      end[divergence_idx + 1 :]     ,      dims[divergence_idx + 1 :]     ,      start_edges=[True for _ in start_edges[divergence_idx + 1 :]]     ,      end_edges=end_edges[divergence_idx + 1 :]     ,      ))

		# If both start and end are at the edges of the subtree rooted at
		# divergence_idx, we can just select the whole subtree at once
		if start_edges[divergence_idx] and end_edges[divergence_idx]:
				slices.append(path + (slice(start[divergence_idx]     ,      end[divergence_idx] + 1),))
		# If just start is at the edge, we can grab almost all of the subtree,
		# treating only the ragged bottom edge as an edge case
		elif start_edges[divergence_idx]:
				slices.append(path + (slice(start[divergence_idx]     ,      end[divergence_idx]),))
				slices.extend(lower())
		# Analogous to the previous case, but the top is ragged this time
		elif end_edges[divergence_idx]:
				slices.extend(upper())
				slices.append(path + (slice(start[divergence_idx] + 1     ,      end[divergence_idx] + 1),))
		# If both sides of the range are ragged, we need to handle both sides
		# separately. If there's contiguous meat in between them, we can index it
		# in one big chunk
		else:
				slices.extend(upper())
				UpperCamelCase__				:					Dict					=					end[divergence_idx] - start[divergence_idx]
				if middle_ground > 1:
						slices.append(path + (slice(start[divergence_idx] + 1     ,      end[divergence_idx]),))
				slices.extend(lower())

		return slices







@torch.jit.ignore
def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_)							->     torch.Tensor:
		UpperCamelCase__				:					List[Any]					=					t.shape[:no_batch_dims]
		UpperCamelCase__				:					Optional[int]					=					list(_flat_idx_to_idx(lowerCamelCase_     ,      lowerCamelCase_))
		# _get_minimal_slice_set is inclusive
		UpperCamelCase__				:					Dict					=					list(_flat_idx_to_idx(flat_end - 1     ,      lowerCamelCase_))

		# Get an ordered list of slices to perform
		UpperCamelCase__				:					int					=					_get_minimal_slice_set(
		    lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      )

		UpperCamelCase__				:					List[Any]					=					[t[s] for s in slices]

		return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors])







def       __UpperCAmelCase						(  lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_     ,      lowerCamelCase_ = False     ,      lowerCamelCase_ = None     ,      lowerCamelCase_ = False     ,      )							->     Any:
		if not (len(lowerCamelCase_) > 0):
				raise ValueError('Must provide at least one input')

		UpperCamelCase__				:					int					=					[shape[:no_batch_dims] for shape in _fetch_dims(lowerCamelCase_)]
		UpperCamelCase__				:					int					=					tuple([max(lowerCamelCase_) for s in zip(*lowerCamelCase_)])

		def _prep_inputs(lowerCamelCase_) -> torch.Tensor:
				if not low_mem:
						if not sum(t.shape[:no_batch_dims]) == no_batch_dims:
								UpperCamelCase__				:					List[Any]					=					t.expand(orig_batch_dims + t.shape[no_batch_dims:])
						UpperCamelCase__				:					Optional[int]					=					t.reshape(-1     ,      *t.shape[no_batch_dims:])
				else:
						UpperCamelCase__				:					Optional[int]					=					t.expand(orig_batch_dims + t.shape[no_batch_dims:])
				return t

		UpperCamelCase__				:					Dict[str, Any]					=					tensor_tree_map(_prep_inputs     ,      lowerCamelCase_)
		UpperCamelCase__				:					int					=					None
		if _out is not None:
				UpperCamelCase__				:					Optional[int]					=					tensor_tree_map(lambda lowerCamelCase_: t.view([-1] + list(t.shape[no_batch_dims:]))     ,      _out)

		UpperCamelCase__				:					Dict					=					1
		for d in orig_batch_dims:
				flat_batch_dim *= d

		UpperCamelCase__				:					int					=					flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

		def _select_chunk(lowerCamelCase_) -> torch.Tensor:
				return t[i : i + chunk_size] if t.shape[0] != 1 else t

		UpperCamelCase__				:					List[Any]					=					0
		UpperCamelCase__				:					Optional[Any]					=					prepped_outputs
		for _ in range(lowerCamelCase_):
				# Chunk the input
				if not low_mem:
						UpperCamelCase__				:					str					=					_select_chunk
				else:
						UpperCamelCase__				:					List[Any]					=					partial(
						    _chunk_slice     ,      flat_start=lowerCamelCase_     ,      flat_end=min(lowerCamelCase_     ,      i + chunk_size)     ,      no_batch_dims=len(lowerCamelCase_)     ,      )

				UpperCamelCase__				:					Dict[str, Any]					=					tensor_tree_map(lowerCamelCase_     ,      lowerCamelCase_)

				# Run the layer on the chunk
				UpperCamelCase__				:					List[Any]					=					layer(**lowerCamelCase_)

				# Allocate space for the output
				if out is None:
						UpperCamelCase__				:					Optional[int]					=					tensor_tree_map(lambda lowerCamelCase_: t.new_zeros((flat_batch_dim,) + t.shape[1:])     ,      lowerCamelCase_)

				# Put the chunk in its pre-allocated space
				if isinstance(lowerCamelCase_     ,      lowerCamelCase_):

						def assign(lowerCamelCase_     ,      lowerCamelCase_) -> None:
								for k, v in da.items():
										if isinstance(lowerCamelCase_     ,      lowerCamelCase_):
												assign(lowerCamelCase_     ,      da[k])
										else:
												if _add_into_out:
														v[i : i + chunk_size] += da[k]
												else:
														UpperCamelCase__				:					List[str]					=					da[k]

						assign(lowerCamelCase_     ,      lowerCamelCase_)
				elif isinstance(lowerCamelCase_     ,      lowerCamelCase_):
						for xa, xa in zip(lowerCamelCase_     ,      lowerCamelCase_):
								if _add_into_out:
										xa[i : i + chunk_size] += xa
								else:
										UpperCamelCase__				:					int					=					xa
				elif isinstance(lowerCamelCase_     ,      torch.Tensor):
						if _add_into_out:
								out[i : i + chunk_size] += output_chunk
						else:
								UpperCamelCase__				:					Dict					=					output_chunk
				else:
						raise ValueError('Not supported')

				i += chunk_size

		UpperCamelCase__				:					int					=					tensor_tree_map(lambda lowerCamelCase_: t.view(orig_batch_dims + t.shape[1:])     ,      lowerCamelCase_)

		return out







class       __lowercase					:

	def __init__(      self   :						List[str]  ,					UpperCAmelCase_   :						int = 512  ,					):
			UpperCamelCase__				:					str					=					max_chunk_size
			UpperCamelCase__				:					Optional[int]					=					None
			UpperCamelCase__				:					Optional[tuple]					=					None

	def      __UpperCamelCase							(      self   :						str  ,					UpperCAmelCase_   :						Callable  ,					UpperCAmelCase_   :						tuple  ,					UpperCAmelCase_   :						int):
			logging.info('Tuning chunk size...')

			if min_chunk_size >= self.max_chunk_size:
					return min_chunk_size

			UpperCamelCase__				:					List[int]					=					[2**l for l in range(int(math.log(self.max_chunk_size  ,					2)) + 1)]
			UpperCamelCase__				:					List[Any]					=					[c for c in candidates if c > min_chunk_size]
			UpperCamelCase__				:					List[Any]					=					[min_chunk_size] + candidates
			candidates[-1] += 4

			def test_chunk_size(UpperCAmelCase_   :						int) -> bool:
					try:
							with torch.no_grad():
									fn(*UpperCAmelCase_  ,					chunk_size=UpperCAmelCase_)
							return True
					except RuntimeError:
							return False

			UpperCamelCase__				:					Tuple					=					0
			UpperCamelCase__				:					Dict					=					len(UpperCAmelCase_) - 1
			while i > min_viable_chunk_size_index:
					UpperCamelCase__				:					Optional[int]					=					test_chunk_size(candidates[i])
					if not viable:
							UpperCamelCase__				:					Tuple					=					(min_viable_chunk_size_index + i) // 2
					else:
							UpperCamelCase__				:					Optional[int]					=					i
							UpperCamelCase__				:					Dict					=					(i + len(UpperCAmelCase_) - 1) // 2

			return candidates[min_viable_chunk_size_index]

	def      __UpperCamelCase							(      self   :						Any  ,					UpperCAmelCase_   :						Iterable  ,					UpperCAmelCase_   :						Iterable):
			UpperCamelCase__				:					List[str]					=					True
			for aa, aa in zip(UpperCAmelCase_  ,					UpperCAmelCase_):
					assert type(UpperCAmelCase_) == type(UpperCAmelCase_)
					if isinstance(UpperCAmelCase_  ,					(list, tuple)):
							consistent &= self._compare_arg_caches(UpperCAmelCase_  ,					UpperCAmelCase_)
					elif isinstance(UpperCAmelCase_  ,					UpperCAmelCase_):
							UpperCamelCase__				:					Union[str, Any]					=					[v for _, v in sorted(aa.items()  ,					key=lambda UpperCAmelCase_: x[0])]
							UpperCamelCase__				:					str					=					[v for _, v in sorted(aa.items()  ,					key=lambda UpperCAmelCase_: x[0])]
							consistent &= self._compare_arg_caches(UpperCAmelCase_  ,					UpperCAmelCase_)
					else:
							consistent &= aa == aa

			return consistent





	def      __UpperCamelCase							(      self   :						List[Any]  ,					UpperCAmelCase_   :						Callable  ,					UpperCAmelCase_   :						tuple  ,					UpperCAmelCase_   :						int  ,					):
			UpperCamelCase__				:					List[Any]					=					True
			UpperCamelCase__				:					tuple					=					tree_map(lambda UpperCAmelCase_: a.shape if isinstance(UpperCAmelCase_  ,					torch.Tensor) else a  ,					UpperCAmelCase_  ,					UpperCAmelCase_)
			if self.cached_arg_data is not None:
					# If args have changed shape/value, we need to re-tune
					assert len(self.cached_arg_data) == len(UpperCAmelCase_)
					UpperCamelCase__				:					Union[str, Any]					=					self._compare_arg_caches(self.cached_arg_data  ,					UpperCAmelCase_)
			else:
					# Otherwise, we can reuse the precomputed value
					UpperCamelCase__				:					Optional[int]					=					False

			if not consistent:
					UpperCamelCase__				:					Tuple					=					self._determine_favorable_chunk_size(
					    UpperCAmelCase_  ,					UpperCAmelCase_  ,					UpperCAmelCase_  ,					)
					UpperCamelCase__				:					Optional[Any]					=					arg_data

			assert self.cached_chunk_size is not None

			return self.cached_chunk_size




