





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
     import os

     # The slow tests are often failing with OOM error on GPU
     # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
     # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
     UpperCamelCase__     =	"platform"

     import jax
     import jax.numpy as jnp

     from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
         FlaxBlenderbotSmallForConditionalGeneration,
         FlaxBlenderbotSmallModel,
         shift_tokens_right,
     )





def       _UpperCamelCase       (a__				:List[Any]				,     a__				:Union[str, Any]				,     a__				:Optional[int]=None				,     a__				:Optional[int]=None				,     a__				:Dict=None				,     a__				:List[str]=None				,     a__				:Dict=None				,     a__				:Optional[Any]=None				,     ):






       """simple docstring"""





       if attention_mask is None:
              UpperCamelCase__			       =     np.where(input_ids != config.pad_token_id				,     1				,     0    )
       if decoder_attention_mask is None:
              UpperCamelCase__			       =     np.where(decoder_input_ids != config.pad_token_id				,     1				,     0    )
       if head_mask is None:
              UpperCamelCase__			       =     np.ones((config.encoder_layers, config.encoder_attention_heads)    )
       if decoder_head_mask is None:
              UpperCamelCase__			       =     np.ones((config.decoder_layers, config.decoder_attention_heads)    )
       if cross_attn_head_mask is None:
              UpperCamelCase__			       =     np.ones((config.decoder_layers, config.decoder_attention_heads)    )
       return {
           "input_ids": input_ids,
           "decoder_input_ids": decoder_input_ids,
           "attention_mask": attention_mask,
           "decoder_attention_mask": attention_mask,
       }

class     __SCREAMING_SNAKE_CASE   :






       def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=13					,	__lowerCAmelCase=7					,	__lowerCAmelCase=True					,	__lowerCAmelCase=False					,	__lowerCAmelCase=99					,	__lowerCAmelCase=16					,	__lowerCAmelCase=2					,	__lowerCAmelCase=4					,	__lowerCAmelCase=4					,	__lowerCAmelCase="gelu"					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=32					,	__lowerCAmelCase=2					,	__lowerCAmelCase=1					,	__lowerCAmelCase=0					,	__lowerCAmelCase=0.02					,	):
              UpperCamelCase__			       =     parent
              UpperCamelCase__			       =     batch_size
              UpperCamelCase__			       =     seq_length
              UpperCamelCase__			       =     is_training
              UpperCamelCase__			       =     use_labels
              UpperCamelCase__			       =     vocab_size
              UpperCamelCase__			       =     hidden_size
              UpperCamelCase__			       =     num_hidden_layers
              UpperCamelCase__			       =     num_attention_heads
              UpperCamelCase__			       =     intermediate_size
              UpperCamelCase__			       =     hidden_act
              UpperCamelCase__			       =     hidden_dropout_prob
              UpperCamelCase__			       =     attention_probs_dropout_prob
              UpperCamelCase__			       =     max_position_embeddings
              UpperCamelCase__			       =     eos_token_id
              UpperCamelCase__			       =     pad_token_id
              UpperCamelCase__			       =     bos_token_id
              UpperCamelCase__			       =     initializer_range






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     np.clip(ids_tensor([self.batch_size, self.seq_length - 1]					,	self.vocab_size					)					,	3					,	self.vocab_size					)
              UpperCamelCase__			       =     np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)					,	dtype=np.intaa					))					,	-1					)

              UpperCamelCase__			       =     shift_tokens_right(__lowerCAmelCase					,	1					,	2					)

              UpperCamelCase__			       =     BlenderbotSmallConfig(
                  vocab_size=self.vocab_size					,	d_model=self.hidden_size					,	encoder_layers=self.num_hidden_layers					,	decoder_layers=self.num_hidden_layers					,	encoder_attention_heads=self.num_attention_heads					,	decoder_attention_heads=self.num_attention_heads					,	encoder_ffn_dim=self.intermediate_size					,	decoder_ffn_dim=self.intermediate_size					,	dropout=self.hidden_dropout_prob					,	attention_dropout=self.attention_probs_dropout_prob					,	max_position_embeddings=self.max_position_embeddings					,	eos_token_id=self.eos_token_id					,	bos_token_id=self.bos_token_id					,	pad_token_id=self.pad_token_id					,	initializer_range=self.initializer_range					,	use_cache=__lowerCAmelCase					,	)
              UpperCamelCase__			       =     prepare_blenderbot_inputs_dict(__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					)
              return config, inputs_dict






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__       ,       UpperCamelCase__			       =     self.prepare_config_and_inputs()
              return config, inputs_dict






       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
              UpperCamelCase__			       =     20
              UpperCamelCase__			       =     model_class_name(__lowerCAmelCase					)

              UpperCamelCase__			       =     model.encode(inputs_dict["""input_ids"""]					)

              UpperCamelCase__       ,       UpperCamelCase__			       =     (
                  inputs_dict["""decoder_input_ids"""],
                  inputs_dict["""decoder_attention_mask"""],
              )

              UpperCamelCase__			       =     model.init_cache(decoder_input_ids.shape[0]					,	__lowerCAmelCase					,	__lowerCAmelCase					)
              UpperCamelCase__			       =     jnp.ones((decoder_input_ids.shape[0], max_decoder_length)					,	dtype="""i4"""					)

              UpperCamelCase__			       =     jnp.broadcast_to(
                  jnp.arange(decoder_input_ids.shape[-1] - 1					)[None, :]					,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)					,	)
              UpperCamelCase__			       =     model.decode(
                  decoder_input_ids[:, :-1]					,	__lowerCAmelCase					,	decoder_attention_mask=__lowerCAmelCase					,	past_key_values=__lowerCAmelCase					,	decoder_position_ids=__lowerCAmelCase					,	)

              UpperCamelCase__			       =     jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]					,	dtype="""i4"""					)
              UpperCamelCase__			       =     model.decode(
                  decoder_input_ids[:, -1:]					,	__lowerCAmelCase					,	decoder_attention_mask=__lowerCAmelCase					,	past_key_values=outputs_cache.past_key_values					,	decoder_position_ids=__lowerCAmelCase					,	)

              UpperCamelCase__			       =     model.decode(__lowerCAmelCase					,	__lowerCAmelCase					)

              UpperCamelCase__			       =     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])					)					)
              self.parent.assertTrue(diff < 1E-3					,	msg=f"""Max diff is {diff}"""					)







       def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
              UpperCamelCase__			       =     20
              UpperCamelCase__			       =     model_class_name(__lowerCAmelCase					)

              UpperCamelCase__			       =     model.encode(inputs_dict["""input_ids"""]					)

              UpperCamelCase__       ,       UpperCamelCase__			       =     (
                  inputs_dict["""decoder_input_ids"""],
                  inputs_dict["""decoder_attention_mask"""],
              )

              UpperCamelCase__			       =     jnp.concatenate(
                  [
                      decoder_attention_mask,
                      jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])					),
                  ]					,	axis=-1					,	)

              UpperCamelCase__			       =     model.init_cache(decoder_input_ids.shape[0]					,	__lowerCAmelCase					,	__lowerCAmelCase					)
              UpperCamelCase__			       =     jnp.broadcast_to(
                  jnp.arange(decoder_input_ids.shape[-1] - 1					)[None, :]					,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)					,	)

              UpperCamelCase__			       =     model.decode(
                  decoder_input_ids[:, :-1]					,	__lowerCAmelCase					,	decoder_attention_mask=__lowerCAmelCase					,	past_key_values=__lowerCAmelCase					,	decoder_position_ids=__lowerCAmelCase					,	)
              UpperCamelCase__			       =     jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]					,	dtype="""i4"""					)
              UpperCamelCase__			       =     model.decode(
                  decoder_input_ids[:, -1:]					,	__lowerCAmelCase					,	past_key_values=outputs_cache.past_key_values					,	decoder_attention_mask=__lowerCAmelCase					,	decoder_position_ids=__lowerCAmelCase					,	)

              UpperCamelCase__			       =     model.decode(__lowerCAmelCase					,	__lowerCAmelCase					,	decoder_attention_mask=__lowerCAmelCase					)

              UpperCamelCase__			       =     np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])					)					)
              self.parent.assertTrue(diff < 1E-3					,	msg=f"""Max diff is {diff}"""					)






@require_flax
class     __SCREAMING_SNAKE_CASE   (       unittest.TestCase    ):
       snake_case       : int         =    99






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     np.array(
                  [
                      [71, 82, 18, 33, 46, 91, 2],
                      [68, 34, 26, 58, 30, 82, 2],
                      [5, 97, 17, 39, 94, 40, 2],
                      [76, 83, 94, 25, 70, 78, 2],
                      [87, 59, 41, 35, 48, 66, 2],
                      [55, 13, 16, 58, 5, 2, 1],  # note padding
                      [64, 27, 31, 51, 12, 75, 2],
                      [52, 64, 86, 17, 83, 39, 2],
                      [48, 61, 9, 24, 71, 82, 2],
                      [26, 1, 60, 48, 22, 13, 2],
                      [21, 5, 62, 28, 14, 76, 2],
                      [45, 98, 37, 86, 59, 48, 2],
                      [70, 70, 50, 9, 28, 0, 2],
                  ]					,	dtype=np.intaa					,	)

              UpperCamelCase__			       =     input_ids.shape[0]
              UpperCamelCase__			       =     BlenderbotSmallConfig(
                  vocab_size=self.vocab_size					,	d_model=24					,	encoder_layers=2					,	decoder_layers=2					,	encoder_attention_heads=2					,	decoder_attention_heads=2					,	encoder_ffn_dim=32					,	decoder_ffn_dim=32					,	max_position_embeddings=48					,	eos_token_id=2					,	pad_token_id=1					,	bos_token_id=0					,	)
              return config, input_ids, batch_size






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__       ,       UpperCamelCase__       ,       UpperCamelCase__			       =     self._get_config_and_data()
              UpperCamelCase__			       =     FlaxBlenderbotSmallForConditionalGeneration(__lowerCAmelCase					)
              UpperCamelCase__			       =     lm_model(input_ids=__lowerCAmelCase					)
              UpperCamelCase__			       =     (batch_size, input_ids.shape[1], config.vocab_size)
              self.assertEqual(outputs["""logits"""].shape					,	__lowerCAmelCase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     BlenderbotSmallConfig(
                  vocab_size=self.vocab_size					,	d_model=14					,	encoder_layers=2					,	decoder_layers=2					,	encoder_attention_heads=2					,	decoder_attention_heads=2					,	encoder_ffn_dim=8					,	decoder_ffn_dim=8					,	max_position_embeddings=48					,	)
              UpperCamelCase__			       =     FlaxBlenderbotSmallForConditionalGeneration(__lowerCAmelCase					)
              UpperCamelCase__			       =     np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]					,	dtype=np.intaa					)
              UpperCamelCase__			       =     np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]					,	dtype=np.intaa					)
              UpperCamelCase__			       =     lm_model(input_ids=__lowerCAmelCase					,	decoder_input_ids=__lowerCAmelCase					)
              UpperCamelCase__			       =     (*summary.shape, config.vocab_size)
              self.assertEqual(outputs["""logits"""].shape					,	__lowerCAmelCase					)







       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]					,	dtype=np.intaa					)
              UpperCamelCase__			       =     shift_tokens_right(__lowerCAmelCase					,	1					,	2					)
              UpperCamelCase__			       =     np.equal(__lowerCAmelCase					,	1					).astype(np.floataa					).sum()
              UpperCamelCase__			       =     np.equal(__lowerCAmelCase					,	1					).astype(np.floataa					).sum()
              self.assertEqual(shifted.shape					,	input_ids.shape					)
              self.assertEqual(__lowerCAmelCase					,	n_pad_before - 1					)
              self.assertTrue(np.equal(shifted[:, 0]					,	2					).all()					)






@require_flax
class     __SCREAMING_SNAKE_CASE   (       _a				,						unittest.TestCase				,						_a    ):
       snake_case       : Union[str, Any]         =    True
       snake_case       : Dict         =    (
           (
               FlaxBlenderbotSmallModel,
               FlaxBlenderbotSmallForConditionalGeneration,
           )
           if is_flax_available()
           else ()
       )
       snake_case       : List[str]         =    (FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__			       =     FlaxBlenderbotSmallModelTester(self					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__       ,       UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              for model_class in self.all_model_classes:
                     self.model_tester.check_use_cache_forward(__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__       ,       UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs()
              for model_class in self.all_model_classes:
                     self.model_tester.check_use_cache_forward_with_attn_mask(__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__       ,       UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     with self.subTest(model_class.__name__					):
                            UpperCamelCase__			       =     self._prepare_for_class(__lowerCAmelCase					,	__lowerCAmelCase					)
                            UpperCamelCase__			       =     model_class(__lowerCAmelCase					)

                            @jax.jit
                            def encode_jitted(__lowerCAmelCase					,	__lowerCAmelCase=None					,	**__lowerCAmelCase					):
                                   return model.encode(input_ids=__lowerCAmelCase					,	attention_mask=__lowerCAmelCase					)

                            with self.subTest("""JIT Enabled"""					):
                                   UpperCamelCase__			       =     encode_jitted(**__lowerCAmelCase					).to_tuple()

                            with self.subTest("""JIT Disabled"""					):
                                   with jax.disable_jit():
                                          UpperCamelCase__			       =     encode_jitted(**__lowerCAmelCase					).to_tuple()

                            self.assertEqual(len(__lowerCAmelCase					)					,	len(__lowerCAmelCase					)					)
                            for jitted_output, output in zip(__lowerCAmelCase					,	__lowerCAmelCase					):
                                   self.assertEqual(jitted_output.shape					,	output.shape					)






       def 			_lowerCamelCase						(  self					):
              UpperCamelCase__       ,       UpperCamelCase__			       =     self.model_tester.prepare_config_and_inputs_for_common()

              for model_class in self.all_model_classes:
                     with self.subTest(model_class.__name__					):
                            UpperCamelCase__			       =     model_class(__lowerCAmelCase					)
                            UpperCamelCase__			       =     model.encode(inputs_dict["""input_ids"""]					,	inputs_dict["""attention_mask"""]					)

                            UpperCamelCase__			       =     {
                                """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
                                """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
                                """encoder_outputs""": encoder_outputs,
                            }

                            @jax.jit
                            def decode_jitted(__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
                                   return model.decode(
                                       decoder_input_ids=__lowerCAmelCase					,	decoder_attention_mask=__lowerCAmelCase					,	encoder_outputs=__lowerCAmelCase					,	)

                            with self.subTest("""JIT Enabled"""					):
                                   UpperCamelCase__			       =     decode_jitted(**__lowerCAmelCase					).to_tuple()

                            with self.subTest("""JIT Disabled"""					):
                                   with jax.disable_jit():
                                          UpperCamelCase__			       =     decode_jitted(**__lowerCAmelCase					).to_tuple()

                            self.assertEqual(len(__lowerCAmelCase					)					,	len(__lowerCAmelCase					)					)
                            for jitted_output, output in zip(__lowerCAmelCase					,	__lowerCAmelCase					):
                                   self.assertEqual(jitted_output.shape					,	output.shape					)







       @slow
       def 			_lowerCamelCase						(  self					):
              for model_class_name in self.all_model_classes:
                     UpperCamelCase__			       =     model_class_name.from_pretrained("""facebook/blenderbot_small-90M"""					)
                     # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
                     UpperCamelCase__			       =     np.ones((1, 1)					) * model.config.eos_token_id
                     UpperCamelCase__			       =     model(__lowerCAmelCase					)
                     self.assertIsNotNone(__lowerCAmelCase					)











from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


UpperCamelCase__     =	[
    "openmmlab/upernet-convnext-tiny",
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
UpperCamelCase__     =	"UperNetConfig"

class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase = 0					,	__lowerCAmelCase = False					,	__lowerCAmelCase = 1					,	):
														super().__init__()
														UpperCamelCase__			       =     nn.Convad(
														    in_channels=__lowerCAmelCase					,	out_channels=__lowerCAmelCase					,	kernel_size=__lowerCAmelCase					,	padding=__lowerCAmelCase					,	bias=__lowerCAmelCase					,	dilation=__lowerCAmelCase					,	)
														UpperCamelCase__			       =     nn.BatchNormad(__lowerCAmelCase					)
														UpperCamelCase__			       =     nn.ReLU()







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     self.conv(__lowerCAmelCase					)
														UpperCamelCase__			       =     self.batch_norm(__lowerCAmelCase					)
														UpperCamelCase__			       =     self.activation(__lowerCAmelCase					)

														return output






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														super().__init__()
														UpperCamelCase__			       =     [
														    nn.AdaptiveAvgPoolad(__lowerCAmelCase					),
														    UperNetConvModule(__lowerCAmelCase					,	__lowerCAmelCase					,	kernel_size=1					),
														]
														for i, layer in enumerate(self.layers					):
																					self.add_module(str(__lowerCAmelCase					)					,	__lowerCAmelCase					)







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     input
														for layer in self.layers:
																					UpperCamelCase__			       =     layer(__lowerCAmelCase					)
														return hidden_state






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														super().__init__()
														UpperCamelCase__			       =     pool_scales
														UpperCamelCase__			       =     align_corners
														UpperCamelCase__			       =     in_channels
														UpperCamelCase__			       =     channels
														UpperCamelCase__			       =     []
														for i, pool_scale in enumerate(__lowerCAmelCase					):
																					UpperCamelCase__			       =     UperNetPyramidPoolingBlock(pool_scale=__lowerCAmelCase					,	in_channels=__lowerCAmelCase					,	channels=__lowerCAmelCase					)
																					self.blocks.append(__lowerCAmelCase					)
																					self.add_module(str(__lowerCAmelCase					)					,	__lowerCAmelCase					)







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     []
														for ppm in self.blocks:
																					UpperCamelCase__			       =     ppm(__lowerCAmelCase					)
																					UpperCamelCase__			       =     nn.functional.interpolate(
																					    __lowerCAmelCase					,	size=x.size()[2:]					,	mode="""bilinear"""					,	align_corners=self.align_corners					)
																					ppm_outs.append(__lowerCAmelCase					)
														return ppm_outs






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					):
														super().__init__()

														UpperCamelCase__			       =     config
														UpperCamelCase__			       =     config.pool_scales  # e.g. (1, 2, 3, 6)
														UpperCamelCase__			       =     in_channels
														UpperCamelCase__			       =     config.hidden_size
														UpperCamelCase__			       =     False
														UpperCamelCase__			       =     nn.Convad(self.channels					,	config.num_labels					,	kernel_size=1					)

														# PSP Module
														UpperCamelCase__			       =     UperNetPyramidPoolingModule(
														    self.pool_scales					,	self.in_channels[-1]					,	self.channels					,	align_corners=self.align_corners					,	)
														UpperCamelCase__			       =     UperNetConvModule(
														    self.in_channels[-1] + len(self.pool_scales					) * self.channels					,	self.channels					,	kernel_size=3					,	padding=1					,	)
														# FPN Module
														UpperCamelCase__			       =     nn.ModuleList()
														UpperCamelCase__			       =     nn.ModuleList()
														for in_channels in self.in_channels[:-1]:  # skip the top layer
																					UpperCamelCase__			       =     UperNetConvModule(__lowerCAmelCase					,	self.channels					,	kernel_size=1					)
																					UpperCamelCase__			       =     UperNetConvModule(self.channels					,	self.channels					,	kernel_size=3					,	padding=1					)
																					self.lateral_convs.append(__lowerCAmelCase					)
																					self.fpn_convs.append(__lowerCAmelCase					)

														UpperCamelCase__			       =     UperNetConvModule(
														    len(self.in_channels					) * self.channels					,	self.channels					,	kernel_size=3					,	padding=1					,	)






							def 			_lowerCamelCase						(  self					):
														self.apply(self._init_weights					)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														if isinstance(__lowerCAmelCase					,	nn.Convad					):
																					module.weight.data.normal_(mean=0.0					,	std=self.config.initializer_range					)
																					if module.bias is not None:
																												module.bias.data.zero_()






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     inputs[-1]
														UpperCamelCase__			       =     [x]
														psp_outs.extend(self.psp_modules(__lowerCAmelCase					)					)
														UpperCamelCase__			       =     torch.cat(__lowerCAmelCase					,	dim=1					)
														UpperCamelCase__			       =     self.bottleneck(__lowerCAmelCase					)

														return output







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														# build laterals
														UpperCamelCase__			       =     [lateral_conv(encoder_hidden_states[i]					) for i, lateral_conv in enumerate(self.lateral_convs					)]

														laterals.append(self.psp_forward(__lowerCAmelCase					)					)

														# build top-down path
														UpperCamelCase__			       =     len(__lowerCAmelCase					)
														for i in range(used_backbone_levels - 1					,	0					,	-1					):
																					UpperCamelCase__			       =     laterals[i - 1].shape[2:]
																					UpperCamelCase__			       =     laterals[i - 1] + nn.functional.interpolate(
																					    laterals[i]					,	size=__lowerCAmelCase					,	mode="""bilinear"""					,	align_corners=self.align_corners					)

														# build outputs
														UpperCamelCase__			       =     [self.fpn_convs[i](laterals[i]					) for i in range(used_backbone_levels - 1					)]
														# append psp feature
														fpn_outs.append(laterals[-1]					)

														for i in range(used_backbone_levels - 1					,	0					,	-1					):
																					UpperCamelCase__			       =     nn.functional.interpolate(
																					    fpn_outs[i]					,	size=fpn_outs[0].shape[2:]					,	mode="""bilinear"""					,	align_corners=self.align_corners					)
														UpperCamelCase__			       =     torch.cat(__lowerCAmelCase					,	dim=1					)
														UpperCamelCase__			       =     self.fpn_bottleneck(__lowerCAmelCase					)
														UpperCamelCase__			       =     self.classifier(__lowerCAmelCase					)

														return output






class     __SCREAMING_SNAKE_CASE   (       nn.Module    ):






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase = 2					,	__lowerCAmelCase = 3					,	__lowerCAmelCase = 1					):
														super().__init__()

														UpperCamelCase__			       =     config
														UpperCamelCase__			       =     config.auxiliary_in_channels
														UpperCamelCase__			       =     config.auxiliary_channels
														UpperCamelCase__			       =     config.auxiliary_num_convs
														UpperCamelCase__			       =     config.auxiliary_concat_input
														UpperCamelCase__			       =     in_index

														UpperCamelCase__			       =     (kernel_size // 2) * dilation
														UpperCamelCase__			       =     []
														convs.append(
														    UperNetConvModule(
														        self.in_channels					,	self.channels					,	kernel_size=__lowerCAmelCase					,	padding=__lowerCAmelCase					,	dilation=__lowerCAmelCase					)					)
														for i in range(self.num_convs - 1					):
																					convs.append(
																					    UperNetConvModule(
																					        self.channels					,	self.channels					,	kernel_size=__lowerCAmelCase					,	padding=__lowerCAmelCase					,	dilation=__lowerCAmelCase					)					)
														if self.num_convs == 0:
																					UpperCamelCase__			       =     nn.Identity()
														else:
																					UpperCamelCase__			       =     nn.Sequential(*__lowerCAmelCase					)
														if self.concat_input:
																					UpperCamelCase__			       =     UperNetConvModule(
																					    self.in_channels + self.channels					,	self.channels					,	kernel_size=__lowerCAmelCase					,	padding=kernel_size // 2					)

														UpperCamelCase__			       =     nn.Convad(self.channels					,	config.num_labels					,	kernel_size=1					)






							def 			_lowerCamelCase						(  self					):
														self.apply(self._init_weights					)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														if isinstance(__lowerCAmelCase					,	nn.Convad					):
																					module.weight.data.normal_(mean=0.0					,	std=self.config.initializer_range					)
																					if module.bias is not None:
																												module.bias.data.zero_()







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														# just take the relevant feature maps
														UpperCamelCase__			       =     encoder_hidden_states[self.in_index]
														UpperCamelCase__			       =     self.convs(__lowerCAmelCase					)
														if self.concat_input:
																					UpperCamelCase__			       =     self.conv_cat(torch.cat([hidden_states, output]					,	dim=1					)					)
														UpperCamelCase__			       =     self.classifier(__lowerCAmelCase					)
														return output






class     __SCREAMING_SNAKE_CASE   (       _a    ):
							snake_case       : Any         =    UperNetConfig
							snake_case       : List[Any]         =    """pixel_values"""
							snake_case       : Optional[Any]         =    True






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														if isinstance(__lowerCAmelCase					,	__lowerCAmelCase					):
																					module.backbone.init_weights()
																					module.decode_head.init_weights()
																					module.auxiliary_head.init_weights()






							def 			_lowerCamelCase						(  self					):
														self.backbone.init_weights()
														self.decode_head.init_weights()
														self.auxiliary_head.init_weights()







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=False					):
														if isinstance(__lowerCAmelCase					,	__lowerCAmelCase					):
																					UpperCamelCase__			       =     value






UpperCamelCase__     =	r"\n    Parameters:\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

UpperCamelCase__     =	r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using\n            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.\n        output_attentions (`bool`, *optional*):\n            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See\n            `attentions` under returned tensors for more detail.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under\n            returned tensors for more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"

@add_start_docstrings(
    """UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes."""				,						_a				,						)
class     __SCREAMING_SNAKE_CASE   (       _a    ):






							def __init__(  self					,	__lowerCAmelCase					):
														super().__init__(__lowerCAmelCase					)

														UpperCamelCase__			       =     AutoBackbone.from_config(config.backbone_config					)

														# Semantic segmentation head(s)
														UpperCamelCase__			       =     UperNetHead(__lowerCAmelCase					,	in_channels=self.backbone.channels					)
														UpperCamelCase__			       =     UperNetFCNHead(__lowerCAmelCase					) if config.use_auxiliary_head else None

														# Initialize weights and apply final processing
														self.post_init()







							@add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format("""batch_size, sequence_length"""					)					)
							@replace_return_docstrings(output_type=__lowerCAmelCase					,	config_class=_CONFIG_FOR_DOC					)
							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	__lowerCAmelCase = None					,	):
														UpperCamelCase__			       =     return_dict if return_dict is not None else self.config.use_return_dict
														UpperCamelCase__			       =     (
														    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
														)
														UpperCamelCase__			       =     output_attentions if output_attentions is not None else self.config.output_attentions

														UpperCamelCase__			       =     self.backbone.forward_with_filtered_kwargs(
														    __lowerCAmelCase					,	output_hidden_states=__lowerCAmelCase					,	output_attentions=__lowerCAmelCase					)
														UpperCamelCase__			       =     outputs.feature_maps

														UpperCamelCase__			       =     self.decode_head(__lowerCAmelCase					)
														UpperCamelCase__			       =     nn.functional.interpolate(__lowerCAmelCase					,	size=pixel_values.shape[2:]					,	mode="""bilinear"""					,	align_corners=__lowerCAmelCase					)

														UpperCamelCase__			       =     None
														if self.auxiliary_head is not None:
																					UpperCamelCase__			       =     self.auxiliary_head(__lowerCAmelCase					)
																					UpperCamelCase__			       =     nn.functional.interpolate(
																					    __lowerCAmelCase					,	size=pixel_values.shape[2:]					,	mode="""bilinear"""					,	align_corners=__lowerCAmelCase					)

														UpperCamelCase__			       =     None
														if labels is not None:
																					if self.config.num_labels == 1:
																												raise ValueError("""The number of labels should be greater than one"""					)
																					else:
																												# compute weighted loss
																												UpperCamelCase__			       =     CrossEntropyLoss(ignore_index=self.config.loss_ignore_index					)
																												UpperCamelCase__			       =     loss_fct(__lowerCAmelCase					,	__lowerCAmelCase					)
																												UpperCamelCase__			       =     loss_fct(__lowerCAmelCase					,	__lowerCAmelCase					)
																												UpperCamelCase__			       =     main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

														if not return_dict:
																					if output_hidden_states:
																												UpperCamelCase__			       =     (logits,) + outputs[1:]
																					else:
																												UpperCamelCase__			       =     (logits,) + outputs[2:]
																					return ((loss,) + output) if loss is not None else output

														return SemanticSegmenterOutput(
														    loss=__lowerCAmelCase					,	logits=__lowerCAmelCase					,	hidden_states=outputs.hidden_states					,	attentions=outputs.attentions					,	)




