





from random import shuffle

import tensorflow as tf
from numpy import array



def 						A_							(	_lowerCAmelCase				,  _lowerCAmelCase	)		->					Optional[Any]:
 UpperCamelCase      :   List[Any]			      =  int(_lowerCAmelCase	)
 assert noofclusters < len(_lowerCAmelCase	)

 # Find out the dimensionality
 UpperCamelCase      :   str			      =  len(vectors[0]	)

 # Will help select random centroids from among the available vectors
 UpperCamelCase      :   List[str]			      =  list(range(len(_lowerCAmelCase	)	)	)
 shuffle(_lowerCAmelCase	)

 # GRAPH OF COMPUTATION
 # We initialize a new graph and set it as the default during each run
 # of this algorithm. This ensures that as this function is called
 # multiple times, the default graph doesn't keep getting crowded with
 # unused ops and Variables from previous function calls.

 UpperCamelCase      :   Any			      =  tf.Graph()

 with graph.as_default():
  # SESSION OF COMPUTATION

  UpperCamelCase      :   Union[str, Any]			      =  tf.Session()

  ##CONSTRUCTING THE ELEMENTS OF COMPUTATION

  ##First lets ensure we have a Variable vector for each centroid,
  ##initialized to one of the vectors from the available data points
  UpperCamelCase      :   Optional[int]			      =  [
      tf.Variable(vectors[vector_indices[i]]	) for i in range(_lowerCAmelCase	)
  ]
  ##These nodes will assign the centroid Variables the appropriate
  ##values
  UpperCamelCase      :   Tuple			      =  tf.placeholder("float64"				,  [dim]	)
  UpperCamelCase      :   str			      =  []
  for centroid in centroids:
   cent_assigns.append(tf.assign(_lowerCAmelCase				,  _lowerCAmelCase	)	)

  ##Variables for cluster assignments of individual vectors(initialized
  ##to 0 at first)
  UpperCamelCase      :   List[str]			      =  [tf.Variable(0	) for i in range(len(_lowerCAmelCase	)	)]
  ##These nodes will assign an assignment Variable the appropriate
  ##value
  UpperCamelCase      :   Any			      =  tf.placeholder("int32"	)
  UpperCamelCase      :   Optional[int]			      =  []
  for assignment in assignments:
   cluster_assigns.append(tf.assign(_lowerCAmelCase				,  _lowerCAmelCase	)	)

  ##Now lets construct the node that will compute the mean
  # The placeholder for the input
  UpperCamelCase      :   Tuple			      =  tf.placeholder("float"				,  [None, dim]	)
  # The Node/op takes the input and computes a mean along the 0th
  # dimension, i.e. the list of input vectors
  UpperCamelCase      :   List[str]			      =  tf.reduce_mean(_lowerCAmelCase				,  0	)

  ##Node for computing Euclidean distances
  # Placeholders for input
  UpperCamelCase      :   Optional[int]			      =  tf.placeholder("float"				,  [dim]	)
  UpperCamelCase      :   Any			      =  tf.placeholder("float"				,  [dim]	)
  UpperCamelCase      :   List[str]			      =  tf.sqrt(tf.reduce_sum(tf.pow(tf.sub(_lowerCAmelCase				,  _lowerCAmelCase	)				,  2	)	)	)

  ##This node will figure out which cluster to assign a vector to,
  ##based on Euclidean distances of the vector from the centroids.
  # Placeholder for input
  UpperCamelCase      :   Tuple			      =  tf.placeholder("float"				,  [noofclusters]	)
  UpperCamelCase      :   str			      =  tf.argmin(_lowerCAmelCase				,  0	)

  ##INITIALIZING STATE VARIABLES

  ##This will help initialization of all Variables defined with respect
  ##to the graph. The Variable-initializer should be defined after
  ##all the Variables have been constructed, so that each of them
  ##will be included in the initialization.
  UpperCamelCase      :   Optional[Any]			      =  tf.initialize_all_variables()

  # Initialize all variables
  sess.run(_lowerCAmelCase	)

  ##CLUSTERING ITERATIONS

  # Now perform the Expectation-Maximization steps of K-Means clustering
  # iterations. To keep things simple, we will only do a set number of
  # iterations, instead of using a Stopping Criterion.
  UpperCamelCase      :   Any			      =  100
  for _ in range(_lowerCAmelCase	):
   ##EXPECTATION STEP
   ##Based on the centroid locations till last iteration, compute
   ##the _expected_ centroid assignments.
   # Iterate over each vector
   for vector_n in range(len(_lowerCAmelCase	)	):
    UpperCamelCase      :   Dict			      =  vectors[vector_n]
    # Compute Euclidean distance between this vector and each
    # centroid. Remember that this list cannot be named
    #'centroid_distances', since that is the input to the
    # cluster assignment node.
    UpperCamelCase      :   Optional[Any]			      =  [
        sess.run(_lowerCAmelCase				,  feed_dict={va: vect, va: sess.run(_lowerCAmelCase	)}	)
        for centroid in centroids
    ]
    # Now use the cluster assignment node, with the distances
    # as the input
    UpperCamelCase      :   Tuple			      =  sess.run(
        _lowerCAmelCase				,  feed_dict={centroid_distances: distances}	)
    # Now assign the value to the appropriate state variable
    sess.run(
        cluster_assigns[vector_n]				,  feed_dict={assignment_value: assignment}	)

   ##MAXIMIZATION STEP
   # Based on the expected state computed from the Expectation Step,
   # compute the locations of the centroids so as to maximize the
   # overall objective of minimizing within-cluster Sum-of-Squares
   for cluster_n in range(_lowerCAmelCase	):
    # Collect all the vectors assigned to this cluster
    UpperCamelCase      :   str			      =  [
        vectors[i]
        for i in range(len(_lowerCAmelCase	)	)
        if sess.run(assignments[i]	) == cluster_n
    ]
    # Compute new centroid location
    UpperCamelCase      :   Optional[int]			      =  sess.run(
        _lowerCAmelCase				,  feed_dict={mean_input: array(_lowerCAmelCase	)}	)
    # Assign value to appropriate variable
    sess.run(
        cent_assigns[cluster_n]				,  feed_dict={centroid_value: new_location}	)

        # Return centroids and assignments
  UpperCamelCase      :   Union[str, Any]			      =  sess.run(_lowerCAmelCase	)
  UpperCamelCase      :   Tuple			      =  sess.run(_lowerCAmelCase	)
  return centroids, assignments







import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
    import torch
    from torch import nn

    from transformers import ViTMAEForPreTraining, ViTMAEModel
    from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor



class    A__  :





      def __init__(					self			,				A_			,				A_=13			,				A_=30			,				A_=2			,				A_=3			,				A_=True			,				A_=True			,				A_=32			,				A_=5			,				A_=4			,				A_=37			,				A_="gelu"			,				A_=0.1			,				A_=0.1			,				A_=10			,				A_=0.02			,				A_=3			,				A_=0.6			,				A_=None			,				):
       '''simple docstring'''




       UpperCamelCase      :   Dict			      =  parent
       UpperCamelCase      :   Union[str, Any]			      =  batch_size
       UpperCamelCase      :   Any			      =  image_size
       UpperCamelCase      :   Tuple			      =  patch_size
       UpperCamelCase      :   Union[str, Any]			      =  num_channels
       UpperCamelCase      :   Tuple			      =  is_training
       UpperCamelCase      :   int			      =  use_labels
       UpperCamelCase      :   Tuple			      =  hidden_size
       UpperCamelCase      :   Optional[Any]			      =  num_hidden_layers
       UpperCamelCase      :   str			      =  num_attention_heads
       UpperCamelCase      :   Optional[int]			      =  intermediate_size
       UpperCamelCase      :   Dict			      =  hidden_act
       UpperCamelCase      :   List[str]			      =  hidden_dropout_prob
       UpperCamelCase      :   str			      =  attention_probs_dropout_prob
       UpperCamelCase      :   Optional[Any]			      =  type_sequence_label_size
       UpperCamelCase      :   Tuple			      =  initializer_range
       UpperCamelCase      :   Any			      =  mask_ratio
       UpperCamelCase      :   int			      =  scope

       # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
       # (we add 1 for the [CLS] token)
       UpperCamelCase      :   Union[str, Any]			      =  (image_size // patch_size) ** 2
       UpperCamelCase      :   int			      =  int(math.ceil((1 - mask_ratio) * (num_patches + 1)  )  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Any			      =  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

       UpperCamelCase      :   int			      =  None
       if self.use_labels:
        UpperCamelCase      :   Optional[Any]			      =  ids_tensor([self.batch_size]			,				self.type_sequence_label_size  )

       UpperCamelCase      :   str			      =  self.get_config()

       return config, pixel_values, labels





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return ViTMAEConfig(
           image_size=self.image_size			,				patch_size=self.patch_size			,				num_channels=self.num_channels			,				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			,				is_decoder=A_			,				initializer_range=self.initializer_range			,				mask_ratio=self.mask_ratio			,				)





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   Any			      =  ViTMAEModel(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Optional[Any]			      =  model(A_  )
       self.parent.assertEqual(result.last_hidden_state.shape			,				(self.batch_size, self.seq_length, self.hidden_size)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   Tuple			      =  ViTMAEForPreTraining(A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   List[str]			      =  model(A_  )
       UpperCamelCase      :   Any			      =  (self.image_size // self.patch_size) ** 2
       UpperCamelCase      :   List[str]			      =  self.patch_size**2 * self.num_channels
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, num_patches, expected_num_channels)  )

       # test greyscale images
       UpperCamelCase      :   str			      =  1
       UpperCamelCase      :   Union[str, Any]			      =  ViTMAEForPreTraining(A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Union[str, Any]			      =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
       UpperCamelCase      :   Union[str, Any]			      =  model(A_  )
       UpperCamelCase      :   List[str]			      =  self.patch_size**2
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, num_patches, expected_num_channels)  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  self.prepare_config_and_inputs()
       UpperCamelCase      ,       UpperCamelCase      ,       UpperCamelCase      :   Optional[int]			      =  config_and_inputs
       UpperCamelCase      :   Optional[int]			      =  {"pixel_values": pixel_values}
       return config, inputs_dict



@require_torch
class    A__  (			__snake_case			,  __snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :Any		       =  (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
      _UpperCAmelCase      :int		       =  {'feature-extraction': ViTMAEModel} if is_torch_available() else {}

      _UpperCAmelCase      :Optional[Any]		       =  False
      _UpperCAmelCase      :Any		       =  False
      _UpperCAmelCase      :Optional[int]		       =  False
      _UpperCAmelCase      :Dict		       =  False





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  ViTMAEModelTester(self  )
       UpperCamelCase      :   Any			      =  ConfigTester(self			,				config_class=A_			,				has_text_modality=A_			,				hidden_size=37  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      @unittest.skip(reason="ViTMAE does not use inputs_embeds"  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      ,       UpperCamelCase      :   List[Any]			      =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase      :   Any			      =  model_class(A_  )
        self.assertIsInstance(model.get_input_embeddings()			,				(nn.Module)  )
        UpperCamelCase      :   Any			      =  model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(A_			,				nn.Linear  )  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      ,       UpperCamelCase      :   Optional[int]			      =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase      :   List[str]			      =  model_class(A_  )
        UpperCamelCase      :   Any			      =  inspect.signature(model.forward  )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        UpperCamelCase      :   List[str]			      =  [*signature.parameters.keys()]

        UpperCamelCase      :   Optional[int]			      =  ["pixel_values"]
        self.assertListEqual(arg_names[:1]			,				A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[int]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_pretraining(*A_  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       np.random.seed(2  )

       UpperCamelCase      :   Union[str, Any]			      =  int((pt_model.config.image_size // pt_model.config.patch_size) ** 2  )
       UpperCamelCase      :   Optional[Any]			      =  np.random.uniform(size=(self.model_tester.batch_size, num_patches)  )
       UpperCamelCase      :   Tuple			      =  torch.from_numpy(A_  )

       # Add `noise` argument.
       # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
       UpperCamelCase      :   Any			      =  pt_noise

       super().check_pt_tf_models(A_			,				A_			,				A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      ,       UpperCamelCase      :   Tuple			      =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase      :   Optional[Any]			      =  model_class(A_  )
        model.to(A_  )
        model.eval()
        # make random mask reproducible
        torch.manual_seed(2  )
        with torch.no_grad():
         UpperCamelCase      :   str			      =  model(**self._prepare_for_class(A_			,				A_  )  )

        UpperCamelCase      :   Any			      =  outputs[0].cpu().numpy()
        UpperCamelCase      :   str			      =  0

        with tempfile.TemporaryDirectory() as tmpdirname:
         model.save_pretrained(A_  )
         UpperCamelCase      :   Optional[int]			      =  model_class.from_pretrained(A_  )
         model.to(A_  )
         # make random mask reproducible
         torch.manual_seed(2  )
         with torch.no_grad():
          UpperCamelCase      :   Tuple			      =  model(**self._prepare_for_class(A_			,				A_  )  )

         # Make sure we don't have nans
         UpperCamelCase      :   Dict			      =  after_outputs[0].cpu().numpy()
         UpperCamelCase      :   int			      =  0
         UpperCamelCase      :   str			      =  np.amax(np.abs(out_a - out_a  )  )
         self.assertLessEqual(A_			,				1e-5  )





      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      @unittest.skip(reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase      :   List[str]			      =  ViTMAEModel.from_pretrained(A_  )
        self.assertIsNotNone(A_  )



def 						A_							(	)		->					Optional[int]:
 UpperCamelCase      :   List[Any]			      =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
 return image



@require_torch
@require_vision
class    A__  (			unittest.TestCase     ):





      @cached_property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"  ) if is_vision_available() else None





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       np.random.seed(2  )

       UpperCamelCase      :   Union[str, Any]			      =  ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"  ).to(A_  )

       UpperCamelCase      :   Any			      =  self.default_image_processor
       UpperCamelCase      :   Optional[Any]			      =  prepare_img()
       UpperCamelCase      :   int			      =  image_processor(images=A_			,				return_tensors="pt"  ).to(A_  )

       # prepare a noise vector that will be also used for testing the TF model
       # (this way we can ensure that the PT and TF models operate on the same inputs)
       UpperCamelCase      :   Dict			      =  ViTMAEConfig()
       UpperCamelCase      :   Any			      =  int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2  )
       UpperCamelCase      :   Optional[Any]			      =  np.random.uniform(size=(1, num_patches)  )

       # forward pass
       with torch.no_grad():
        UpperCamelCase      :   Optional[Any]			      =  model(**A_			,				noise=torch.from_numpy(A_  ).to(device=A_  )  )

       # verify the logits
       UpperCamelCase      :   str			      =  torch.Size((1, 196, 768)  )
       self.assertEqual(outputs.logits.shape			,				A_  )

       UpperCamelCase      :   str			      =  torch.tensor(
           [[-0.05_48, -1.70_23, -0.93_25], [0.37_21, -0.56_70, -0.22_33], [0.82_35, -1.38_78, -0.35_24]]  )

       self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]			,				expected_slice.to(A_  )			,				atol=1e-4  )  )
