'''simple docstring'''






import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


SCREAMING_SNAKE_CASE_	   =			logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE_	   =			'PoolFormerConfig'

# Base docstring
SCREAMING_SNAKE_CASE_	   =			'sail/poolformer_s12'
SCREAMING_SNAKE_CASE_	   =			[1, 5_12, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE_	   =			'sail/poolformer_s12'
SCREAMING_SNAKE_CASE_	   =			'tabby, tabby cat'

SCREAMING_SNAKE_CASE_	   =			[
    'sail/poolformer_s12',
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]






def  UpperCamelCase__					(					_lowercase    :				int   ,    _lowercase    :				Dict = 0.0   ,    _lowercase    :				Any = False     )							->	Any:
		if drop_prob == 0.0 or not training:
				return input
		__UpperCAmelCase:					List[str]		= 1 - drop_prob
		__UpperCAmelCase:					Optional[int]		= (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
		__UpperCAmelCase:					str		= keep_prob + torch.rand(snake_case_   ,    dtype=input.dtype   ,    device=input.device     )
		random_tensor.floor_()  # binarize
		__UpperCAmelCase:					List[str]		= input.div(snake_case_     ) * random_tensor
		return output




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_ = None					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					List[str]		= drop_prob







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				return drop_path(_a   ,    self.drop_prob   ,    self.training					)







		def    lowercase_  (							self					):

				'''simple docstring'''







				return "p={}".format(self.drop_prob					)




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_=None					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					int		= patch_size if isinstance(_a   ,    collections.abc.Iterable					) else (patch_size, patch_size)
				__UpperCAmelCase:					Optional[Any]		= stride if isinstance(_a   ,    collections.abc.Iterable					) else (stride, stride)
				__UpperCAmelCase:					Dict		= padding if isinstance(_a   ,    collections.abc.Iterable					) else (padding, padding)

				__UpperCAmelCase:					Optional[Any]		= nn.Convad(_a   ,    _a   ,    kernel_size=_a   ,    stride=_a   ,    padding=_a					)
				__UpperCAmelCase:					str		= norm_layer(_a					) if norm_layer else nn.Identity()







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					Any		= self.projection(_a					)
				__UpperCAmelCase:					List[Any]		= self.norm(_a					)
				return embeddings




class    a (       nn.GroupNorm				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    **snake_case_					):

				'''simple docstring'''







				super().__init__(1   ,    _a   ,    **_a					)




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					Union[str, Any]		= nn.AvgPoolad(_a   ,    stride=1   ,    padding=pool_size // 2   ,    count_include_pad=_a					)







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				return self.pool(_a					) - hidden_states




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					Tuple		= nn.Convad(_a   ,    _a   ,    1					)
				__UpperCAmelCase:					Tuple		= nn.Convad(_a   ,    _a   ,    1					)
				__UpperCAmelCase:					Any		= PoolFormerDropPath(_a					)
				if isinstance(config.hidden_act   ,    _a					):
						__UpperCAmelCase:					Optional[int]		= ACTaFN[config.hidden_act]
				else:
						__UpperCAmelCase:					Optional[Any]		= config.hidden_act







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					Tuple		= self.conva(_a					)
				__UpperCAmelCase:					int		= self.act_fn(_a					)
				__UpperCAmelCase:					List[str]		= self.drop(_a					)
				__UpperCAmelCase:					Union[str, Any]		= self.conva(_a					)
				__UpperCAmelCase:					Optional[int]		= self.drop(_a					)

				return hidden_states




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					List[str]		= PoolFormerPooling(_a					)
				__UpperCAmelCase:					int		= PoolFormerOutput(_a   ,    _a   ,    _a   ,    _a					)
				__UpperCAmelCase:					Optional[Any]		= PoolFormerGroupNorm(_a					)
				__UpperCAmelCase:					Optional[Any]		= PoolFormerGroupNorm(_a					)

				# Useful for training neural nets
				__UpperCAmelCase:					List[Any]		= PoolFormerDropPath(_a					) if drop_path > 0.0 else nn.Identity()
				__UpperCAmelCase:					List[str]		= config.use_layer_scale
				if config.use_layer_scale:
						__UpperCAmelCase:					Tuple		= nn.Parameter(
						    config.layer_scale_init_value * torch.ones((_a)					)   ,    requires_grad=_a					)
						__UpperCAmelCase:					Tuple		= nn.Parameter(
						    config.layer_scale_init_value * torch.ones((_a)					)   ,    requires_grad=_a					)







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				if self.use_layer_scale:
						__UpperCAmelCase:					Optional[Any]		= self.pooling(self.before_norm(_a					)					)
						__UpperCAmelCase:					Union[str, Any]		= self.layer_scale_a.unsqueeze(-1					).unsqueeze(-1					) * pooling_output
						# First residual connection
						__UpperCAmelCase:					Union[str, Any]		= hidden_states + self.drop_path(_a					)
						__UpperCAmelCase:					List[Any]		= ()

						__UpperCAmelCase:					List[str]		= self.output(self.after_norm(_a					)					)
						__UpperCAmelCase:					Optional[Any]		= self.layer_scale_a.unsqueeze(-1					).unsqueeze(-1					) * layer_output
						# Second residual connection
						__UpperCAmelCase:					List[Any]		= hidden_states + self.drop_path(_a					)

						__UpperCAmelCase:					Union[str, Any]		= (output,) + outputs
						return outputs

				else:
						__UpperCAmelCase:					int		= self.drop_path(self.pooling(self.before_norm(_a					)					)					)
						# First residual connection
						__UpperCAmelCase:					int		= pooling_output + hidden_states
						__UpperCAmelCase:					Dict		= ()

						# Second residual connection inside the PoolFormerOutput block
						__UpperCAmelCase:					int		= self.drop_path(self.output(self.after_norm(_a					)					)					)
						__UpperCAmelCase:					Union[str, Any]		= hidden_states + layer_output

						__UpperCAmelCase:					Dict		= (output,) + outputs
						return outputs




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					Union[str, Any]		= config
				# stochastic depth decay rule
				__UpperCAmelCase:					Union[str, Any]		= [x.item() for x in torch.linspace(0   ,    config.drop_path_rate   ,    sum(config.depths					)					)]

				# patch embeddings
				__UpperCAmelCase:					Optional[Any]		= []
				for i in range(config.num_encoder_blocks					):
						embeddings.append(
						    PoolFormerEmbeddings(
						        patch_size=config.patch_sizes[i]   ,    stride=config.strides[i]   ,    padding=config.padding[i]   ,    num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]   ,    hidden_size=config.hidden_sizes[i]   ,    )					)
				__UpperCAmelCase:					Optional[int]		= nn.ModuleList(_a					)

				# Transformer blocks
				__UpperCAmelCase:					Optional[Any]		= []
				__UpperCAmelCase:					Tuple		= 0
				for i in range(config.num_encoder_blocks					):
						# each block consists of layers
						__UpperCAmelCase:					List[Any]		= []
						if i != 0:
								cur += config.depths[i - 1]
						for j in range(config.depths[i]					):
								layers.append(
								    PoolFormerLayer(
								        _a   ,    num_channels=config.hidden_sizes[i]   ,    pool_size=config.pool_size   ,    hidden_size=config.hidden_sizes[i]   ,    intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio					)   ,    drop_path=dpr[cur + j]   ,    )					)
						blocks.append(nn.ModuleList(_a					)					)

				__UpperCAmelCase:					List[Any]		= nn.ModuleList(_a					)







		def    lowercase_  (							self   ,    snake_case_   ,    snake_case_=False   ,    snake_case_=True					):

				'''simple docstring'''







				__UpperCAmelCase:					Tuple		= () if output_hidden_states else None

				__UpperCAmelCase:					List[str]		= pixel_values
				for idx, layers in enumerate(zip(self.patch_embeddings   ,    self.block					)					):
						__UpperCAmelCase:					Union[str, Any]		= layers
						# Get patch embeddings from hidden_states
						__UpperCAmelCase:					str		= embedding_layer(_a					)
						# Send the embeddings through the blocks
						for _, blk in enumerate(_a					):
								__UpperCAmelCase:					Optional[Any]		= blk(_a					)
								__UpperCAmelCase:					List[str]		= layer_outputs[0]

						if output_hidden_states:
								__UpperCAmelCase:					Any		= all_hidden_states + (hidden_states,)

				if not return_dict:
						return tuple(v for v in [hidden_states, all_hidden_states] if v is not None					)

				return BaseModelOutputWithNoAttention(last_hidden_state=_a   ,    hidden_states=_a					)




class    a (       UpperCamelCase__				):



		"""simple docstring"""
		__lowerCAmelCase				      =					PoolFormerConfig
		__lowerCAmelCase				      =					"""poolformer"""
		__lowerCAmelCase				      =					"""pixel_values"""
		__lowerCAmelCase				      =					True







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				if isinstance(_a   ,    (nn.Linear, nn.Convad)					):
						module.weight.data.normal_(mean=0.0   ,    std=self.config.initializer_range					)
						if module.bias is not None:
								module.bias.data.zero_()
				elif isinstance(_a   ,    nn.LayerNorm					):
						module.bias.data.zero_()
						module.weight.data.fill_(1.0					)







		def    lowercase_  (							self   ,    snake_case_   ,    snake_case_=False					):

				'''simple docstring'''







				if isinstance(_a   ,    _a					):
						__UpperCAmelCase:					List[str]		= value



SCREAMING_SNAKE_CASE_	   =			R'\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\n    Parameters:\n        config ([`PoolFormerConfig`]): 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'

SCREAMING_SNAKE_CASE_	   =			R'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`PoolFormerImageProcessor.__call__`] for details.\n'




@add_start_docstrings(
    """The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."""						,  UpperCamelCase__						,  )
class    a (       UpperCamelCase__				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__(_a					)
				__UpperCAmelCase:					Any		= config

				__UpperCAmelCase:					Union[str, Any]		= PoolFormerEncoder(_a					)

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







		def    lowercase_  (							self					):

				'''simple docstring'''







				return self.embeddings.patch_embeddings







		@add_start_docstrings_to_model_forward(_a					)
		@add_code_sample_docstrings(
		    checkpoint=_CHECKPOINT_FOR_DOC   ,    output_type=_a   ,    config_class=_CONFIG_FOR_DOC   ,    modality="""vision"""   ,    expected_output=_EXPECTED_OUTPUT_SHAPE   ,    )
		def    lowercase_  (							self   ,    snake_case_ = None   ,    snake_case_ = None   ,    snake_case_ = None   ,    ):

				'''simple docstring'''







				__UpperCAmelCase:					Optional[int]		= (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				__UpperCAmelCase:					List[str]		= return_dict if return_dict is not None else self.config.use_return_dict

				if pixel_values is None:
						raise ValueError("""You have to specify pixel_values"""					)

				__UpperCAmelCase:					Any		= self.encoder(
				    _a   ,    output_hidden_states=_a   ,    return_dict=_a   ,    )
				__UpperCAmelCase:					Optional[int]		= encoder_outputs[0]

				if not return_dict:
						return (sequence_output, None) + encoder_outputs[1:]

				return BaseModelOutputWithNoAttention(
				    last_hidden_state=_a   ,    hidden_states=encoder_outputs.hidden_states   ,    )




class    a (       nn.Module				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__()
				__UpperCAmelCase:					List[str]		= nn.Linear(config.hidden_size   ,    config.hidden_size					)







		def    lowercase_  (							self   ,    snake_case_					):

				'''simple docstring'''







				__UpperCAmelCase:					str		= self.dense(_a					)
				return output




@add_start_docstrings(
    """\n    PoolFormer Model transformer with an image classification head on top\n    """						,  UpperCamelCase__						,  )
class    a (       UpperCamelCase__				):



		"""simple docstring"""







		def __init__(							self   ,    snake_case_					):

				'''simple docstring'''







				super().__init__(_a					)
				__UpperCAmelCase:					Tuple		= config.num_labels
				__UpperCAmelCase:					Optional[Any]		= PoolFormerModel(_a					)

				# Final norm
				__UpperCAmelCase:					Optional[int]		= PoolFormerGroupNorm(config.hidden_sizes[-1]					)
				# Classifier head
				__UpperCAmelCase:					str		= (
				    nn.Linear(config.hidden_sizes[-1]   ,    config.num_labels					) if config.num_labels > 0 else nn.Identity()
				)

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







		@add_start_docstrings_to_model_forward(_a					)
		@add_code_sample_docstrings(
		    checkpoint=_IMAGE_CLASS_CHECKPOINT   ,    output_type=_a   ,    config_class=_CONFIG_FOR_DOC   ,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,    )
		def    lowercase_  (							self   ,    snake_case_ = None   ,    snake_case_ = None   ,    snake_case_ = None   ,    snake_case_ = None   ,    ):

				'''simple docstring'''







				__UpperCAmelCase:					Optional[Any]		= return_dict if return_dict is not None else self.config.use_return_dict

				__UpperCAmelCase:					Optional[Any]		= self.poolformer(
				    _a   ,    output_hidden_states=_a   ,    return_dict=_a   ,    )

				__UpperCAmelCase:					Tuple		= outputs[0]

				__UpperCAmelCase:					Any		= self.classifier(self.norm(_a					).mean([-2, -1]					)					)

				__UpperCAmelCase:					Dict		= None
				if labels is not None:
						if self.config.problem_type is None:
								if self.num_labels == 1:
										__UpperCAmelCase:					Union[str, Any]		= """regression"""
								elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
										__UpperCAmelCase:					Optional[int]		= """single_label_classification"""
								else:
										__UpperCAmelCase:					List[Any]		= """multi_label_classification"""

						if self.config.problem_type == "regression":
								__UpperCAmelCase:					Union[str, Any]		= MSELoss()
								if self.num_labels == 1:
										__UpperCAmelCase:					str		= loss_fct(logits.squeeze()   ,    labels.squeeze()					)
								else:
										__UpperCAmelCase:					Optional[Any]		= loss_fct(_a   ,    _a					)
						elif self.config.problem_type == "single_label_classification":
								__UpperCAmelCase:					List[str]		= CrossEntropyLoss()
								__UpperCAmelCase:					Any		= loss_fct(logits.view(-1   ,    self.num_labels					)   ,    labels.view(-1					)					)
						elif self.config.problem_type == "multi_label_classification":
								__UpperCAmelCase:					List[Any]		= BCEWithLogitsLoss()
								__UpperCAmelCase:					Dict		= loss_fct(_a   ,    _a					)

				if not return_dict:
						__UpperCAmelCase:					List[str]		= (logits,) + outputs[2:]
						return ((loss,) + output) if loss is not None else output

				return ImageClassifierOutputWithNoAttention(loss=_a   ,    logits=_a   ,    hidden_states=outputs.hidden_states					)
'''simple docstring'''






import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
    import jax.numpy as jnp

    from transformers.models.albert.modeling_flax_albert import (
        FlaxAlbertForMaskedLM,
        FlaxAlbertForMultipleChoice,
        FlaxAlbertForPreTraining,
        FlaxAlbertForQuestionAnswering,
        FlaxAlbertForSequenceClassification,
        FlaxAlbertForTokenClassification,
        FlaxAlbertModel,
    )




class    a (       unittest.TestCase				):



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    snake_case_=13   ,    snake_case_=7   ,    snake_case_=True   ,    snake_case_=True   ,    snake_case_=True   ,    snake_case_=True   ,    snake_case_=99   ,    snake_case_=32   ,    snake_case_=5   ,    snake_case_=4   ,    snake_case_=37   ,    snake_case_="gelu"   ,    snake_case_=0.1   ,    snake_case_=0.1   ,    snake_case_=512   ,    snake_case_=16   ,    snake_case_=2   ,    snake_case_=0.0_2   ,    snake_case_=4   ,    ):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[Any]		= parent
    __UpperCAmelCase:					List[str]		= batch_size
    __UpperCAmelCase:					Optional[int]		= seq_length
    __UpperCAmelCase:					Optional[Any]		= is_training
    __UpperCAmelCase:					Any		= use_attention_mask
    __UpperCAmelCase:					List[str]		= use_token_type_ids
    __UpperCAmelCase:					List[str]		= use_labels
    __UpperCAmelCase:					List[str]		= vocab_size
    __UpperCAmelCase:					Optional[Any]		= hidden_size
    __UpperCAmelCase:					List[Any]		= num_hidden_layers
    __UpperCAmelCase:					List[Any]		= num_attention_heads
    __UpperCAmelCase:					Tuple		= intermediate_size
    __UpperCAmelCase:					Dict		= hidden_act
    __UpperCAmelCase:					Dict		= hidden_dropout_prob
    __UpperCAmelCase:					Tuple		= attention_probs_dropout_prob
    __UpperCAmelCase:					List[str]		= max_position_embeddings
    __UpperCAmelCase:					List[Any]		= type_vocab_size
    __UpperCAmelCase:					List[Any]		= type_sequence_label_size
    __UpperCAmelCase:					List[str]		= initializer_range
    __UpperCAmelCase:					List[Any]		= num_choices







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					int		= ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size					)

    __UpperCAmelCase:					str		= None
    if self.use_attention_mask:
      __UpperCAmelCase:					Any		= random_attention_mask([self.batch_size, self.seq_length]					)

    __UpperCAmelCase:					List[Any]		= None
    if self.use_token_type_ids:
      __UpperCAmelCase:					Optional[int]		= ids_tensor([self.batch_size, self.seq_length]   ,    self.type_vocab_size					)

    __UpperCAmelCase:					List[Any]		= AlbertConfig(
        vocab_size=self.vocab_size   ,    hidden_size=self.hidden_size   ,    num_hidden_layers=self.num_hidden_layers   ,    num_attention_heads=self.num_attention_heads   ,    intermediate_size=self.intermediate_size   ,    hidden_act=self.hidden_act   ,    hidden_dropout_prob=self.hidden_dropout_prob   ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,    max_position_embeddings=self.max_position_embeddings   ,    type_vocab_size=self.type_vocab_size   ,    is_decoder=snake_case_   ,    initializer_range=self.initializer_range   ,    )

    return config, input_ids, token_type_ids, attention_mask







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Tuple		= self.prepare_config_and_inputs()
    __UpperCAmelCase, __UpperCAmelCase, __UpperCAmelCase, __UpperCAmelCase:					List[str]		= config_and_inputs
    __UpperCAmelCase:					List[Any]		= {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
    return config, inputs_dict




@require_flax
class    a (       __lowerCAmelCase						,  unittest.TestCase				):



  """simple docstring"""
  __lowerCAmelCase				      =					(
      (
          FlaxAlbertModel,
          FlaxAlbertForPreTraining,
          FlaxAlbertForMaskedLM,
          FlaxAlbertForMultipleChoice,
          FlaxAlbertForQuestionAnswering,
          FlaxAlbertForSequenceClassification,
          FlaxAlbertForTokenClassification,
          FlaxAlbertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Union[str, Any]		= FlaxAlbertModelTester(self					)







  @slow
  def    lowercase_  (							self					):

    '''simple docstring'''







    for model_class_name in self.all_model_classes:
      __UpperCAmelCase:					str		= model_class_name.from_pretrained("""albert-base-v2"""					)
      __UpperCAmelCase:					int		= model(np.ones((1, 1)					)					)
      self.assertIsNotNone(snake_case_					)




@require_flax
class    a (       unittest.TestCase				):



  """simple docstring"""







  @slow
  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					str		= FlaxAlbertModel.from_pretrained("""albert-base-v2"""					)
    __UpperCAmelCase:					Dict		= np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]					)
    __UpperCAmelCase:					int		= np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
    __UpperCAmelCase:					Optional[int]		= model(snake_case_   ,    attention_mask=snake_case_					)[0]
    __UpperCAmelCase:					Optional[int]		= (1, 11, 768)
    self.assertEqual(output.shape   ,    snake_case_					)
    __UpperCAmelCase:					str		= np.array(
        [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]					)

    self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]   ,    snake_case_   ,    atol=1e-4					)					)