


'''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 			_a			(		_lowercase      :			List[str]						):





     '''simple docstring'''






     __UpperCAmelCase			:				Optional[int]													=    []
     if isinstance(lowercase_    ,      lowercase_						):
          for v in tree.values():
               shapes.extend(_fetch_dims(lowercase_						)						)
     elif isinstance(lowercase_    ,      (list, tuple)						):
          for t in tree:
               shapes.extend(_fetch_dims(lowercase_						)						)
     elif isinstance(lowercase_    ,      torch.Tensor						):
          shapes.append(tree.shape						)
     else:
          raise ValueError('''Not supported'''						)

     return shapes




@torch.jit.ignore
def 			_a			(		_lowercase      :			int    ,      _lowercase      :			Optional[int]						):





     '''simple docstring'''






     __UpperCAmelCase			:				Any													=    []
     for d in reversed(lowercase_						):
          idx.append(flat_idx % d						)
          __UpperCAmelCase			:				Dict													=    flat_idx // d

     return tuple(reversed(lowercase_						)						)




@torch.jit.ignore
def 			_a			(		_lowercase      :			List[str]    ,      _lowercase      :			int    ,      _lowercase      :			Optional[Any]    ,      _lowercase      :			List[str] = None    ,      _lowercase      :			int = None    ,      ):





     '''simple docstring'''






     def reduce_edge_list(_lowercase      :			Any						) -> None:
          __UpperCAmelCase			:				List[Any]													=    True
          for i in range(len(lowercase_						)						):
               __UpperCAmelCase			:				Dict													=    -1 * (i + 1)
               l[reversed_idx] &= tally
               __UpperCAmelCase			:				str													=    l[reversed_idx]

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

     # Base cases. Either start/end are empty and we're done, or the final,
     # one-dimensional tensor can be simply sliced
     if len(lowercase_						) == 0:
          return [()]
     elif len(lowercase_						) == 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(lowercase_    ,      lowercase_						):
          if s == e:
               path_list.append(slice(lowercase_    ,      s + 1						)						)
          else:
               break

     __UpperCAmelCase			:				Tuple[slice, ...]													=    tuple(lowercase_						)
     __UpperCAmelCase			:				List[Any]													=    len(lowercase_						)

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

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

          __UpperCAmelCase			:				int													=    start[divergence_idx]
          return tuple(
              path + (slice(lowercase_    ,      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[Any]													=    end[divergence_idx]
          return tuple(
              path + (slice(lowercase_    ,      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			:				Any													=    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 			_a			(		_lowercase      :			List[str]    ,      _lowercase      :			str    ,      _lowercase      :			Optional[int]    ,      _lowercase      :			Union[str, Any]						):





     '''simple docstring'''






     __UpperCAmelCase			:				Any													=    t.shape[:no_batch_dims]
     __UpperCAmelCase			:				List[Any]													=    list(_flat_idx_to_idx(lowercase_    ,      lowercase_						)						)
     # _get_minimal_slice_set is inclusive
     __UpperCAmelCase			:				Tuple													=    list(_flat_idx_to_idx(flat_end - 1    ,      lowercase_						)						)

     # Get an ordered list of slices to perform
     __UpperCAmelCase			:				Dict													=    _get_minimal_slice_set(
         lowercase_    ,      lowercase_    ,      lowercase_    ,      )

     __UpperCAmelCase			:				Optional[Any]													=    [t[s] for s in slices]

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




def 			_a			(		_lowercase      :			Dict    ,      _lowercase      :			Dict    ,      _lowercase      :			List[Any]    ,      _lowercase      :			Optional[int]    ,      _lowercase      :			int = False    ,      _lowercase      :			int = None    ,      _lowercase      :			List[str] = False    ,      ):





     '''simple docstring'''






     if not (len(lowercase_						) > 0):
          raise ValueError('''Must provide at least one input'''						)

     __UpperCAmelCase			:				int													=    [shape[:no_batch_dims] for shape in _fetch_dims(lowercase_						)]
     __UpperCAmelCase			:				List[Any]													=    tuple([max(lowercase_						) for s in zip(*lowercase_						)]						)

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

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

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

     __UpperCAmelCase			:				Dict													=    flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

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

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

          __UpperCAmelCase			:				Dict[str, Any]													=    tensor_tree_map(lowercase_    ,      lowercase_						)

          # Run the layer on the chunk
          __UpperCAmelCase			:				List[str]													=    layer(**lowercase_						)

          # Allocate space for the output
          if out is None:
               __UpperCAmelCase			:				Optional[Any]													=    tensor_tree_map(lambda _lowercase						: t.new_zeros((flat_batch_dim,) + t.shape[1:]						)    ,      lowercase_						)

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

               def assign(_lowercase      :			Optional[Any]    ,      _lowercase      :			List[str]						) -> None:
                    for k, v in da.items():
                         if isinstance(lowercase_    ,      lowercase_						):
                              assign(lowercase_    ,      da[k]						)
                         else:
                              if _add_into_out:
                                   v[i : i + chunk_size] += da[k]
                              else:
                                   __UpperCAmelCase			:				int													=    da[k]

               assign(lowercase_    ,      lowercase_						)
          elif isinstance(lowercase_    ,      lowercase_						):
               for xa, xa in zip(lowercase_    ,      lowercase_						):
                    if _add_into_out:
                         xa[i : i + chunk_size] += xa
                    else:
                         __UpperCAmelCase			:				str													=    xa
          elif isinstance(lowercase_    ,      torch.Tensor						):
               if _add_into_out:
                    out[i : i + chunk_size] += output_chunk
               else:
                    __UpperCAmelCase			:				Union[str, Any]													=    output_chunk
          else:
               raise ValueError('''Not supported'''						)

          i += chunk_size

     __UpperCAmelCase			:				Optional[Any]													=    tensor_tree_map(lambda _lowercase						: t.view(orig_batch_dims + t.shape[1:]						)    ,      lowercase_						)

     return out



class 			_a       :


     """simple docstring"""






     def __init__(       self  :	Optional[int]							,				snake_case  :	int = 512							,				)		->							Optional[Any]:
          __UpperCAmelCase			:				Optional[Any]													=    max_chunk_size
          __UpperCAmelCase			:				Optional[int]													=    None
          __UpperCAmelCase			:				Optional[tuple]													=    None

     def 				lowerCamelCase__				(       self  :	List[Any]							,				snake_case  :	Callable							,				snake_case  :	tuple							,				snake_case  :	int    )		->							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			:				Union[str, 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(snake_case  :	int    ) -> bool:
               try:
                    with torch.no_grad():
                         fn(*__a							,				chunk_size=__a    )
                    return True
               except RuntimeError:
                    return False

          __UpperCAmelCase			:				Tuple													=    0
          __UpperCAmelCase			:				Tuple													=    len(__a    ) - 1
          while i > min_viable_chunk_size_index:
               __UpperCAmelCase			:				int													=    test_chunk_size(candidates[i]    )
               if not viable:
                    __UpperCAmelCase			:				Dict													=    (min_viable_chunk_size_index + i) // 2
               else:
                    __UpperCAmelCase			:				List[str]													=    i
                    __UpperCAmelCase			:				int													=    (i + len(__a    ) - 1) // 2

          return candidates[min_viable_chunk_size_index]

     def 				lowerCamelCase__				(       self  :	List[Any]							,				snake_case  :	Iterable							,				snake_case  :	Iterable    )		->							bool:
          __UpperCAmelCase			:				Optional[Any]													=    True
          for aa, aa in zip(__a							,				__a    ):
               assert type(__a    ) == type(__a    )
               if isinstance(__a							,				(list, tuple)    ):
                    consistent &= self._compare_arg_caches(__a							,				__a    )
               elif isinstance(__a							,				__a    ):
                    __UpperCAmelCase			:				List[Any]													=    [v for _, v in sorted(aa.items()							,				key=lambda snake_case    : x[0]    )]
                    __UpperCAmelCase			:				Tuple													=    [v for _, v in sorted(aa.items()							,				key=lambda snake_case    : x[0]    )]
                    consistent &= self._compare_arg_caches(__a							,				__a    )
               else:
                    consistent &= aa == aa

          return consistent






     def 				lowerCamelCase__				(       self  :	str							,				snake_case  :	Callable							,				snake_case  :	tuple							,				snake_case  :	int							,				)		->							int:
          __UpperCAmelCase			:				Dict													=    True
          __UpperCAmelCase			:				tuple													=    tree_map(lambda snake_case    : a.shape if isinstance(__a							,				torch.Tensor    ) else a							,				__a							,				__a    )
          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(__a    )
               __UpperCAmelCase			:				Any													=    self._compare_arg_caches(self.cached_arg_data							,				__a    )
          else:
               # Otherwise, we can reuse the precomputed value
               __UpperCAmelCase			:				Optional[int]													=    False

          if not consistent:
               __UpperCAmelCase			:				Optional[Any]													=    self._determine_favorable_chunk_size(
                   __a							,				__a							,				__a							,				)
               __UpperCAmelCase			:				int													=    arg_data

          assert self.cached_chunk_size is not None

          return self.cached_chunk_size



'''simple docstring'''



import unittest

from transformers import AutoTokenizer, NystromformerConfig, is_torch_available
from transformers.testing_utils import require_torch, 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 (
          NystromformerForMaskedLM,
          NystromformerForMultipleChoice,
          NystromformerForQuestionAnswering,
          NystromformerForSequenceClassification,
          NystromformerForTokenClassification,
          NystromformerModel,
      )
      from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST



class 			a       :


     """simple docstring"""






     def __init__(       self  :	Union[str, Any]							,				snake_case  :	str							,				snake_case  :	Union[str, Any]=13							,				snake_case  :	int=7							,				snake_case  :	Union[str, Any]=True							,				snake_case  :	Tuple=True							,				snake_case  :	Tuple=True							,				snake_case  :	Union[str, Any]=True							,				snake_case  :	Tuple=99							,				snake_case  :	str=32							,				snake_case  :	Tuple=5							,				snake_case  :	int=4							,				snake_case  :	Union[str, Any]=37							,				snake_case  :	Union[str, Any]="gelu"							,				snake_case  :	List[Any]=0.1							,				snake_case  :	List[Any]=0.1							,				snake_case  :	Any=512							,				snake_case  :	Dict=16							,				snake_case  :	Dict=2							,				snake_case  :	int=0.02							,				snake_case  :	Dict=3							,				snake_case  :	Tuple=4							,				snake_case  :	Any=None							,				)		->							Tuple:
          __UpperCAmelCase			:				Optional[Any]													=    parent
          __UpperCAmelCase			:				int													=    batch_size
          __UpperCAmelCase			:				List[Any]													=    seq_length
          __UpperCAmelCase			:				Dict													=    is_training
          __UpperCAmelCase			:				Union[str, Any]													=    use_input_mask
          __UpperCAmelCase			:				Union[str, Any]													=    use_token_type_ids
          __UpperCAmelCase			:				Optional[Any]													=    use_labels
          __UpperCAmelCase			:				str													=    vocab_size
          __UpperCAmelCase			:				List[str]													=    hidden_size
          __UpperCAmelCase			:				Any													=    num_hidden_layers
          __UpperCAmelCase			:				Optional[int]													=    num_attention_heads
          __UpperCAmelCase			:				str													=    intermediate_size
          __UpperCAmelCase			:				Tuple													=    hidden_act
          __UpperCAmelCase			:				Tuple													=    hidden_dropout_prob
          __UpperCAmelCase			:				Optional[Any]													=    attention_probs_dropout_prob
          __UpperCAmelCase			:				List[str]													=    max_position_embeddings
          __UpperCAmelCase			:				Optional[Any]													=    type_vocab_size
          __UpperCAmelCase			:				Dict													=    type_sequence_label_size
          __UpperCAmelCase			:				Tuple													=    initializer_range
          __UpperCAmelCase			:				Union[str, Any]													=    num_labels
          __UpperCAmelCase			:				List[str]													=    num_choices
          __UpperCAmelCase			:				List[str]													=    scope

     def 				lowerCamelCase__				(       self  :	Optional[int]    )		->							Any:
          __UpperCAmelCase			:				int													=    ids_tensor([self.batch_size, self.seq_length]							,				self.vocab_size    )

          __UpperCAmelCase			:				Optional[int]													=    None
          if self.use_input_mask:
               __UpperCAmelCase			:				Tuple													=    random_attention_mask([self.batch_size, self.seq_length]    )

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

          __UpperCAmelCase			:				Union[str, Any]													=    None
          __UpperCAmelCase			:				int													=    None
          __UpperCAmelCase			:				Union[str, Any]													=    None
          if self.use_labels:
               __UpperCAmelCase			:				List[Any]													=    ids_tensor([self.batch_size]							,				self.type_sequence_label_size    )
               __UpperCAmelCase			:				int													=    ids_tensor([self.batch_size, self.seq_length]							,				self.num_labels    )
               __UpperCAmelCase			:				int													=    ids_tensor([self.batch_size]							,				self.num_choices    )

          __UpperCAmelCase			:				int													=    self.get_config()

          return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

     def 				lowerCamelCase__				(       self  :	Union[str, Any]    )		->							int:
          return NystromformerConfig(
              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							,				)

     def 				lowerCamelCase__				(       self  :	int							,				snake_case  :	Optional[int]							,				snake_case  :	str							,				snake_case  :	Optional[Any]							,				snake_case  :	List[str]							,				snake_case  :	str							,				snake_case  :	Union[str, Any]							,				snake_case  :	Tuple    )		->							Dict:
          __UpperCAmelCase			:				Any													=    NystromformerModel(config=snake_case    )
          model.to(snake_case    )
          model.eval()
          __UpperCAmelCase			:				str													=    model(snake_case							,				attention_mask=snake_case							,				token_type_ids=snake_case    )
          __UpperCAmelCase			:				List[str]													=    model(snake_case							,				token_type_ids=snake_case    )
          __UpperCAmelCase			:				Dict													=    model(snake_case    )
          self.parent.assertEqual(result.last_hidden_state.shape							,				(self.batch_size, self.seq_length, self.hidden_size)    )

     def 				lowerCamelCase__				(       self  :	Dict							,				snake_case  :	int							,				snake_case  :	Optional[int]							,				snake_case  :	int							,				snake_case  :	int							,				snake_case  :	List[str]							,				snake_case  :	Any							,				snake_case  :	Tuple    )		->							List[Any]:
          __UpperCAmelCase			:				List[Any]													=    NystromformerForMaskedLM(config=snake_case    )
          model.to(snake_case    )
          model.eval()
          __UpperCAmelCase			:				List[str]													=    model(snake_case							,				attention_mask=snake_case							,				token_type_ids=snake_case							,				labels=snake_case    )
          self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.seq_length, self.vocab_size)    )

     def 				lowerCamelCase__				(       self  :	str							,				snake_case  :	Optional[Any]							,				snake_case  :	Union[str, Any]							,				snake_case  :	List[str]							,				snake_case  :	Optional[Any]							,				snake_case  :	Union[str, Any]							,				snake_case  :	List[Any]							,				snake_case  :	Tuple    )		->							Union[str, Any]:
          __UpperCAmelCase			:				List[str]													=    NystromformerForQuestionAnswering(config=snake_case    )
          model.to(snake_case    )
          model.eval()
          __UpperCAmelCase			:				List[str]													=    model(
              snake_case							,				attention_mask=snake_case							,				token_type_ids=snake_case							,				start_positions=snake_case							,				end_positions=snake_case							,				)
          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 				lowerCamelCase__				(       self  :	Dict							,				snake_case  :	Any							,				snake_case  :	int							,				snake_case  :	Dict							,				snake_case  :	Any							,				snake_case  :	Optional[int]							,				snake_case  :	Tuple							,				snake_case  :	Any    )		->							List[str]:
          __UpperCAmelCase			:				int													=    self.num_labels
          __UpperCAmelCase			:				str													=    NystromformerForSequenceClassification(snake_case    )
          model.to(snake_case    )
          model.eval()
          __UpperCAmelCase			:				List[str]													=    model(snake_case							,				attention_mask=snake_case							,				token_type_ids=snake_case							,				labels=snake_case    )
          self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.num_labels)    )

     def 				lowerCamelCase__				(       self  :	Tuple							,				snake_case  :	List[Any]							,				snake_case  :	Optional[int]							,				snake_case  :	int							,				snake_case  :	List[str]							,				snake_case  :	Union[str, Any]							,				snake_case  :	List[str]							,				snake_case  :	List[str]    )		->							Optional[Any]:
          __UpperCAmelCase			:				str													=    self.num_labels
          __UpperCAmelCase			:				int													=    NystromformerForTokenClassification(config=snake_case    )
          model.to(snake_case    )
          model.eval()
          __UpperCAmelCase			:				str													=    model(snake_case							,				attention_mask=snake_case							,				token_type_ids=snake_case							,				labels=snake_case    )
          self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.seq_length, self.num_labels)    )

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






     def 				lowerCamelCase__				(       self  :	List[Any]    )		->							Optional[int]:
          __UpperCAmelCase			:				Union[str, Any]													=    self.prepare_config_and_inputs()
          (
              (
              __UpperCAmelCase
          )							,     (
              __UpperCAmelCase
          )							,     (
              __UpperCAmelCase
          )							,     (
              __UpperCAmelCase
          )							,     (
              __UpperCAmelCase
          )							,     (
              __UpperCAmelCase
          )							,     (
              __UpperCAmelCase
          )							,     
          )			:				int													=    config_and_inputs
          __UpperCAmelCase			:				Tuple													=    {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
          return config, inputs_dict



@require_torch
class 			a       (	_a			,						_a			,						unittest.TestCase ):


     """simple docstring"""






     SCREAMING_SNAKE_CASE :				List[str]							     =     (
         (
             NystromformerModel,
             NystromformerForMaskedLM,
             NystromformerForMultipleChoice,
             NystromformerForQuestionAnswering,
             NystromformerForSequenceClassification,
             NystromformerForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     SCREAMING_SNAKE_CASE :				Union[str, Any]							     =     (
         {
             "feature-extraction": NystromformerModel,
             "fill-mask": NystromformerForMaskedLM,
             "question-answering": NystromformerForQuestionAnswering,
             "text-classification": NystromformerForSequenceClassification,
             "token-classification": NystromformerForTokenClassification,
             "zero-shot": NystromformerForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     SCREAMING_SNAKE_CASE :				Tuple							     =     False
     SCREAMING_SNAKE_CASE :				Union[str, Any]							     =     False

     def 				lowerCamelCase__				(       self  :	Dict    )		->							List[str]:
          __UpperCAmelCase			:				str													=    NystromformerModelTester(self    )
          __UpperCAmelCase			:				List[str]													=    ConfigTester(self							,				config_class=snake_case							,				hidden_size=37    )

     def 				lowerCamelCase__				(       self  :	int    )		->							Any:
          self.config_tester.run_common_tests()

     def 				lowerCamelCase__				(       self  :	Union[str, Any]    )		->							Optional[int]:
          __UpperCAmelCase			:				List[str]													=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case    )

     def 				lowerCamelCase__				(       self  :	Dict    )		->							Tuple:
          __UpperCAmelCase			:				Dict													=    self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
               __UpperCAmelCase			:				List[str]													=    type
               self.model_tester.create_and_check_model(*snake_case    )

     def 				lowerCamelCase__				(       self  :	Any    )		->							List[Any]:
          __UpperCAmelCase			:				List[Any]													=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*snake_case    )

     def 				lowerCamelCase__				(       self  :	Optional[Any]    )		->							Optional[Any]:
          __UpperCAmelCase			:				Optional[Any]													=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_multiple_choice(*snake_case    )

     def 				lowerCamelCase__				(       self  :	str    )		->							Optional[Any]:
          __UpperCAmelCase			:				int													=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*snake_case    )

     def 				lowerCamelCase__				(       self  :	str    )		->							Union[str, Any]:
          __UpperCAmelCase			:				Any													=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*snake_case    )

     def 				lowerCamelCase__				(       self  :	str    )		->							int:
          __UpperCAmelCase			:				Dict													=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*snake_case    )






     @slow
     def 				lowerCamelCase__				(       self  :	Any    )		->							Union[str, Any]:
          for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __UpperCAmelCase			:				List[Any]													=    NystromformerModel.from_pretrained(snake_case    )
               self.assertIsNotNone(snake_case    )



@require_torch
class 			a       (	unittest.TestCase ):


     """simple docstring"""






     @slow
     def 				lowerCamelCase__				(       self  :	List[str]    )		->							List[str]:
          __UpperCAmelCase			:				Union[str, Any]													=    NystromformerModel.from_pretrained('''uw-madison/nystromformer-512'''    )
          __UpperCAmelCase			:				Any													=    torch.tensor([[0, 1, 2, 3, 4, 5]]    )

          with torch.no_grad():
               __UpperCAmelCase			:				Union[str, Any]													=    model(snake_case    )[0]

          __UpperCAmelCase			:				Tuple													=    torch.Size((1, 6, 768)    )
          self.assertEqual(output.shape							,				snake_case    )

          __UpperCAmelCase			:				Tuple													=    torch.tensor(
              [[[-0.4_532, -0.0_936, 0.5_137], [-0.2_676, 0.0_628, 0.6_186], [-0.3_629, -0.1_726, 0.4_716]]]    )

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






     @slow
     def 				lowerCamelCase__				(       self  :	Optional[int]    )		->							str:
          __UpperCAmelCase			:				str													=    '''the [MASK] of Belgium is Brussels'''

          __UpperCAmelCase			:				Dict													=    AutoTokenizer.from_pretrained('''uw-madison/nystromformer-512'''    )
          __UpperCAmelCase			:				Tuple													=    NystromformerForMaskedLM.from_pretrained('''uw-madison/nystromformer-512'''    )

          __UpperCAmelCase			:				Optional[Any]													=    tokenizer(snake_case							,				return_tensors='''pt'''    )

          with torch.no_grad():
               __UpperCAmelCase			:				Union[str, Any]													=    model(encoding.input_ids    ).logits

          __UpperCAmelCase			:				str													=    token_logits[:, 2, :].argmax(-1    )[0]

          self.assertEqual(tokenizer.decode(snake_case    )							,				'''capital'''    )