codingutf8 2022 facebook ai and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tf 2 0 vit mae masked autoencoder model from future import annotations import collections abc import math from copy import deepcopy from dataclasses import dataclass from typing import optional tuple union import numpy as np import tensorflow as tf from activationstf import gettfactivation from fileutils import modeloutput addstartdocstrings addstartdocstringstomodelforward replacereturndocstrings from modelingtfoutputs import tfbasemodeloutput from modelingtfutils import tfmodelinputtype tfpretrainedmodel getinitializer kerasserializable unpackinputs from tfutils import shapelist stablesoftmax from utils import logging from configurationvitmae import vitmaeconfig logger logging getloggername configfordoc vitmaeconfig checkpointfordoc facebookvitmaebase dataclass class tfvitmaemodeloutputmodeloutput lasthiddenstate tf tensor none mask tf tensor none idsrestore tf tensor none hiddenstates tupletf tensor none none attentions tupletf tensor none none dataclass class tfvitmaedecoderoutputmodeloutput logits tf tensor none hiddenstates tupletf tensor none none attentions tupletf tensor none none dataclass class tfvitmaeforpretrainingoutputmodeloutput loss tf tensor none none logits tf tensor none mask tf tensor none idsrestore tf tensor none hiddenstates tupletf tensor none none attentions tupletf tensor none none def get2dsincosposembedembeddim gridsize addclstokenfalse gridh tf rangegridsize dtypetf float32 gridw tf rangegridsize dtypetf float32 grid tf meshgridgridw gridh here w goes first grid tf stackgrid axis0 grid tf reshapegrid 2 1 gridsize gridsize posembed get2dsincosposembedfromgridembeddim grid if addclstoken posembed tf concattf zeros1 embeddim posembed axis0 return posembed def get2dsincosposembedfromgridembeddim grid if embeddim 2 0 raise valueerrorembeddim must be even use half of dimensions to encode gridh embh get1dsincosposembedfromgridembeddim 2 grid0 hw d2 embw get1dsincosposembedfromgridembeddim 2 grid1 hw d2 emb tf concatembh embw axis1 hw d return emb def get1dsincosposembedfromgridembeddim pos if embeddim 2 0 raise valueerrorembeddim must be even omega tf rangeembeddim 2 dtypefloat32 omega embeddim 2 0 omega 1 0 10000omega d2 pos tf reshapepos 1 m out tf einsumm dmd pos omega m d2 outer product half of the positions get sinusoidal pattern and the rest gets cosine pattern and then they are concatenated embsin tf sinout m d2 embcos tf cosout m d2 emb tf concatembsin embcos axis1 m d return emb class tfvitmaeembeddingstf keras layers layer def initself config vitmaeconfig kwargs super initkwargs self patchembeddings tfvitmaepatchembeddingsconfig namepatchembeddings self numpatches self patchembeddings numpatches self config config def buildself inputshape tf tensorshape self clstoken self addweight shape1 1 self config hiddensize initializertf randomnormalinitializerstddevself config initializerrange trainabletrue nameclstoken self positionembeddings self addweight shape1 self numpatches 1 self config hiddensize initializerzeros trainablefalse fixed sincos embedding namepositionembeddings posembed get2dsincosposembed self positionembeddings shape1 intself patchembeddings numpatches0 5 addclstokentrue none self positionembeddings assignposembed super buildinputshape def randommaskingself sequence tf tensor noise tf tensor none none batchsize seqlength dim shapelistsequence lenkeep intseqlength 1 self config maskratio if noise is none noise tf random uniformshapebatchsize seqlength minval0 0 maxval1 0 noise in 0 1 sort noise for each sample idsshuffle tf argsortnoise axis1 ascend small is keep large is remove idsrestore tf argsortidsshuffle axis1 keep the first subset idskeep idsshuffle lenkeep sequenceunmasked tf gather sequence axis1 batchdims1 indicesidskeep generate the binary mask 0 is keep 1 is remove this hack is needed because tf s eagertensors don t support assignment maskkeep tf zerosbatchsize lenkeep maskremove tf onesbatchsize seqlength lenkeep mask tf concatmaskkeep maskremove axis1 unshuffle to get the binary mask mask tf gathermask axis1 batchdims1 indicesidsrestore return sequenceunmasked mask idsrestore def callself pixelvalues tf tensor noise tf tensor none tf tensor embeddings self patchembeddingspixelvalues add position embeddings wo cls token embeddings embeddings self positionembeddings 1 masking length length config maskratio embeddings mask idsrestore self randommaskingembeddings noise append cls token clstoken self clstoken self positionembeddings 1 clstokens tf tileclstoken shapelistembeddings0 1 1 embeddings tf concatclstokens embeddings axis1 return embeddings mask idsrestore class tfvitmaepatchembeddingstf keras layers layer def initself config vitmaeconfig kwargs super initkwargs imagesize patchsize config imagesize config patchsize numchannels hiddensize config numchannels config hiddensize imagesize imagesize if isinstanceimagesize collections abc iterable else imagesize imagesize patchsize patchsize if isinstancepatchsize collections abc iterable else patchsize patchsize numpatches imagesize1 patchsize1 imagesize0 patchsize0 self imagesize imagesize self patchsize patchsize self numpatches numpatches self numchannels numchannels self config config self projection tf keras layers conv2d filtershiddensize kernelsizepatchsize stridespatchsize paddingvalid dataformatchannelslast kernelinitializerglorotuniform following torch nn linear biasinitializerzeros nameprojection def callself pixelvalues tf tensor training bool false tf tensor batchsize numchannels height width shapelistpixelvalues if tf executingeagerly if numchannels self numchannels raise valueerror make sure that the channel dimension of the pixel values match with the one set in the configuration if height self imagesize0 or width self imagesize1 raise valueerror finput image size heightwidth doesn t match model f self imagesize0self imagesize1 when running on cpu tf keras layers conv2d doesn t support nchw format so change the input format from nchw to nhwc shape batchsize inheight inwidth inchannelsnumchannels pixelvalues tf transposepixelvalues perm0 2 3 1 projection self projectionpixelvalues change the 2d spatial dimensions to a single temporal dimension shape batchsize numpatches outchannelsembeddim numpatches width self patchsize1 height self patchsize0 x tf reshapetensorprojection shapebatchsize numpatches 1 return x copied from transformers models vit modelingtfvit tfvitselfattention with vitvitmae class tfvitmaeselfattentiontf keras layers layer def initself config vitmaeconfig kwargs super initkwargs if config hiddensize config numattentionheads 0 raise valueerror fthe hidden size config hiddensize is not a multiple of the number fof attention heads config numattentionheads self numattentionheads config numattentionheads self attentionheadsize intconfig hiddensize config numattentionheads self allheadsize self numattentionheads self attentionheadsize self sqrtattheadsize math sqrtself attentionheadsize self query tf keras layers dense unitsself allheadsize kernelinitializergetinitializerconfig initializerrange namequery self key tf keras layers dense unitsself allheadsize kernelinitializergetinitializerconfig initializerrange namekey self value tf keras layers dense unitsself allheadsize kernelinitializergetinitializerconfig initializerrange namevalue self dropout tf keras layers dropoutrateconfig attentionprobsdropoutprob def transposeforscoresself tensor tf tensor batchsize int tf tensor reshape from batchsize seqlength allheadsize to batchsize seqlength numattentionheads attentionheadsize tensor tf reshapetensortensor shapebatchsize 1 self numattentionheads self attentionheadsize transpose the tensor from batchsize seqlength numattentionheads attentionheadsize to batchsize numattentionheads seqlength attentionheadsize return tf transposetensor perm0 2 1 3 def call self hiddenstates tf tensor headmask tf tensor outputattentions bool training bool false tupletf tensor batchsize shapelisthiddenstates0 mixedquerylayer self queryinputshiddenstates mixedkeylayer self keyinputshiddenstates mixedvaluelayer self valueinputshiddenstates querylayer self transposeforscoresmixedquerylayer batchsize keylayer self transposeforscoresmixedkeylayer batchsize valuelayer self transposeforscoresmixedvaluelayer batchsize take the dot product between query and key to get the raw attention scores batch size numheads seqlenq seqlenk attentionscores tf matmulquerylayer keylayer transposebtrue dk tf castself sqrtattheadsize dtypeattentionscores dtype attentionscores tf divideattentionscores dk normalize the attention scores to probabilities attentionprobs stablesoftmaxlogitsattentionscores axis1 this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper attentionprobs self dropoutinputsattentionprobs trainingtraining mask heads if we want to if headmask is not none attentionprobs tf multiplyattentionprobs headmask attentionoutput tf matmulattentionprobs valuelayer attentionoutput tf transposeattentionoutput perm0 2 1 3 batchsize seqlenq allheadsize attentionoutput tf reshapetensorattentionoutput shapebatchsize 1 self allheadsize outputs attentionoutput attentionprobs if outputattentions else attentionoutput return outputs copied from transformers models vit modelingtfvit tfvitselfoutput with vitvitmae class tfvitmaeselfoutputtf keras layers layer def initself config vitmaeconfig kwargs super initkwargs self dense tf keras layers dense unitsconfig hiddensize kernelinitializergetinitializerconfig initializerrange namedense self dropout tf keras layers dropoutrateconfig hiddendropoutprob def callself hiddenstates tf tensor inputtensor tf tensor training bool false tf tensor hiddenstates self denseinputshiddenstates hiddenstates self dropoutinputshiddenstates trainingtraining return hiddenstates copied from transformers models vit modelingtfvit tfvitattention with vitvitmae class tfvitmaeattentiontf keras layers layer def initself config vitmaeconfig kwargs super initkwargs self selfattention tfvitmaeselfattentionconfig nameattention self denseoutput tfvitmaeselfoutputconfig nameoutput def pruneheadsself heads raise notimplementederror def call self inputtensor tf tensor headmask tf tensor outputattentions bool training bool false tupletf tensor selfoutputs self selfattention hiddenstatesinputtensor headmaskheadmask outputattentionsoutputattentions trainingtraining attentionoutput self denseoutput hiddenstatesselfoutputs0 inputtensorinputtensor trainingtraining outputs attentionoutput selfoutputs1 add attentions if we output them return outputs copied from transformers models vit modelingtfvit tfvitintermediate with vitvitmae class tfvitmaeintermediatetf keras layers layer def initself config vitmaeconfig kwargs super initkwargs self dense tf keras layers dense unitsconfig intermediatesize kernelinitializergetinitializerconfig initializerrange namedense if isinstanceconfig hiddenact str self intermediateactfn gettfactivationconfig hiddenact else self intermediateactfn config hiddenact def callself hiddenstates tf tensor tf tensor hiddenstates self denseinputshiddenstates hiddenstates self intermediateactfnhiddenstates return hiddenstates copied from transformers models vit modelingtfvit tfvitoutput with vitvitmae class tfvitmaeoutputtf keras layers layer def initself config vitmaeconfig kwargs super initkwargs self dense tf keras layers dense unitsconfig hiddensize kernelinitializergetinitializerconfig initializerrange namedense self dropout tf keras layers dropoutrateconfig hiddendropoutprob def callself hiddenstates tf tensor inputtensor tf tensor training bool false tf tensor hiddenstates self denseinputshiddenstates hiddenstates self dropoutinputshiddenstates trainingtraining hiddenstates hiddenstates inputtensor return hiddenstates copied from transformers models vit modelingtfvit tfvitlayer with vitvitmae class tfvitmaelayertf keras layers layer in vitmae layernorm is applied before selfattention first residual connection in vitmae layernorm is also applied after selfattention second residual connection is done here copied from transformers models vit modelingtfvit tfvitencoder with vitvitmae add last layer prunes heads of the model headstoprune dict of layernum list of heads to prune in this layer see base class pretrainedmodel prepare head mask if needed 1 0 in headmask indicate we keep the head attentionprobs has shape bsz x nheads x n x n input headmask has shape numheads or numhiddenlayers x numheads and headmask is converted to shape numhiddenlayers x batch x numheads x seqlength x seqlength an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models vitmaeinputsdocstring r args pixelvalues np ndarray tf tensor listtf tensor dictstr tf tensor or dictstr np ndarray and each example must have the shape batchsize numchannels height width pixel values pixel values can be obtained using autoimageprocessor see vitimageprocessor call for details headmask np ndarray or tf tensor of shape numheads or numlayers numheads optional mask to nullify selected heads of the selfattention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead returndict bool optional whether or not to return a fileutils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation embed tokens append mask tokens to sequence add pos embed apply transformer layers blocks predictor projection remove cls token args pixelvalues tf tensor of shape batchsize height width numchannels or batchsize numchannels height width pixel values returns tf tensor of shape batchsize numpatches patchsize2 numchannels patchified pixel values make sure channels are last sanity checks patchify args patchifiedpixelvalues tf tensor of shape batchsize numpatches patchsize2 numchannels patchified pixel values returns tf tensor of shape batchsize height width numchannels pixel values sanity check unpatchify args pixelvalues tf tensor of shape batchsize height width numchannels pixel values pred tf tensor of shape batchsize numpatches patchsize2 numchannels predicted pixel values mask tf tensor of shape batchsize sequencelength tensor indicating which patches are masked 1 and which are not 0 returns tf tensor pixel reconstruction loss coding utf 8 2022 facebook ai and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license tf 2 0 vit mae masked autoencoder model class for tfvitmaemodel s outputs with potential hidden states and attentions args last_hidden_state tf tensor of shape batch_size sequence_length hidden_size sequence of hidden states at the output of the last layer of the model mask tf tensor of shape batch_size sequence_length tensor indicating which patches are masked 1 and which are not 0 ids_restore tf tensor of shape batch_size sequence_length tensor containing the original index of the shuffled masked patches hidden_states tuple tf tensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple tf tensor optional returned when output_attentions true is passed or when config output_attentions true tuple of tf tensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads class for tfvitmaedecoder s outputs with potential hidden states and attentions args logits tf tensor of shape batch_size sequence_length patch_size 2 num_channels pixel reconstruction logits hidden_states tuple tf tensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple tf tensor optional returned when output_attentions true is passed or when config output_attentions true tuple of tf tensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads class for tfvitmaeforpretraining s outputs with potential hidden states and attentions args loss tf tensor of shape 1 pixel reconstruction loss logits tf tensor of shape batch_size sequence_length patch_size 2 num_channels pixel reconstruction logits mask tf tensor of shape batch_size sequence_length tensor indicating which patches are masked 1 and which are not 0 ids_restore tf tensor of shape batch_size sequence_length tensor containing the original index of the shuffled masked patches hidden_states tuple tf tensor optional returned when output_hidden_states true is passed or when config output_hidden_states true tuple of tf tensor one for the output of the embeddings one for the output of each layer of shape batch_size sequence_length hidden_size hidden states of the model at the output of each layer plus the initial embedding outputs attentions tuple tf tensor optional returned when output_attentions true is passed or when config output_attentions true tuple of tf tensor one for each layer of shape batch_size num_heads sequence_length sequence_length attentions weights after the attention softmax used to compute the weighted average in the self attention heads create 2d sin cos positional embeddings args embed_dim int embedding dimension grid_size int the grid height and width add_cls_token bool optional defaults to false whether or not to add a classification cls token returns tf tensor of shape grid_size grid_size embed_dim or 1 grid_size grid_size embed_dim the position embeddings with or without classification token here w goes first use half of dimensions to encode grid_h h w d 2 h w d 2 h w d embed_dim output dimension for each position pos a list of positions to be encoded size m out m d d 2 m m d 2 outer product half of the positions get sinusoidal pattern and the rest gets cosine pattern and then they are concatenated m d 2 m d 2 m d construct the cls token position and patch embeddings fixed sin cos embedding perform per sample random masking by per sample shuffling per sample shuffling is done by argsort random noise args sequence tf tensor of shape batch_size sequence_length dim noise tf tensor of shape batch_size sequence_length optional which is mainly used for testing purposes to control randomness and maintain the reproducibility noise in 0 1 sort noise for each sample ascend small is keep large is remove keep the first subset generate the binary mask 0 is keep 1 is remove this hack is needed because tf s eagertensors don t support assignment unshuffle to get the binary mask add position embeddings w o cls token masking length length config mask_ratio append cls token this class turns pixel_values of shape batch_size num_channels height width into the initial hidden_states patch embeddings of shape batch_size seq_length hidden_size to be consumed by a transformer following torch nn linear when running on cpu tf keras layers conv2d doesn t support nchw format so change the input format from nchw to nhwc shape batch_size in_height in_width in_channels num_channels change the 2d spatial dimensions to a single temporal dimension shape batch_size num_patches out_channels embed_dim copied from transformers models vit modeling_tf_vit tfvitselfattention with vit vitmae reshape from batch_size seq_length all_head_size to batch_size seq_length num_attention_heads attention_head_size transpose the tensor from batch_size seq_length num_attention_heads attention_head_size to batch_size num_attention_heads seq_length attention_head_size take the dot product between query and key to get the raw attention scores batch size num_heads seq_len_q seq_len_k normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper mask heads if we want to batch_size seq_len_q all_head_size copied from transformers models vit modeling_tf_vit tfvitselfoutput with vit vitmae the residual connection is defined in tfvitmaelayer instead of here as is the case with other models due to the layernorm applied before each block copied from transformers models vit modeling_tf_vit tfvitattention with vit vitmae add attentions if we output them copied from transformers models vit modeling_tf_vit tfvitintermediate with vit vitmae copied from transformers models vit modeling_tf_vit tfvitoutput with vit vitmae copied from transformers models vit modeling_tf_vit tfvitlayer with vit vitmae this corresponds to the block class in the timm implementation in vitmae layernorm is applied before self attention first residual connection in vitmae layernorm is also applied after self attention second residual connection is done here add attentions if we output them copied from transformers models vit modeling_tf_vit tfvitencoder with vit vitmae add last layer prunes heads of the model heads_to_prune dict of layer_num list of heads to prune in this layer see base class pretrainedmodel prepare head mask if needed 1 0 in head_mask indicate we keep the head attention_probs has shape bsz x n_heads x n x n input head_mask has shape num_heads or num_hidden_layers x num_heads and head_mask is converted to shape num_hidden_layers x batch x num_heads x seq_length x seq_length an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models this model inherits from tfpretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a tf keras model https www tensorflow org api_docs python tf keras model subclass use it as a regular tf 2 0 keras model and refer to the tf 2 0 documentation for all matter related to general usage and behavior tip tensorflow models and layers in transformers accept two formats as input having all inputs as keyword arguments like pytorch models or having all inputs as a list tuple or dict in the first positional argument the reason the second format is supported is that keras methods prefer this format when passing inputs to models and layers because of this support when using methods like model fit things should just work for you just pass your inputs and labels in any format that model fit supports if however you want to use the second format outside of keras methods like fit and predict such as when creating your own layers or models with the keras functional api there are three possibilities you can use to gather all the input tensors in the first positional argument a single tensor with pixel_values only and nothing else model pixel_values a list of varying length with one or several input tensors in the order given in the docstring model pixel_values attention_mask or model pixel_values attention_mask token_type_ids a dictionary with one or several input tensors associated to the input names given in the docstring model pixel_values pixel_values token_type_ids token_type_ids note that when creating models and layers with subclassing https keras io guides making_new_layers_and_models_via_subclassing then you don t need to worry about any of this as you can just pass inputs like you would to any other python function tip args config vitmaeconfig model configuration class with all the parameters of the model initializing with a config file does not load the weights associated with the model only the configuration check out the tfpretrainedmodel from_pretrained method to load the model weights args pixel_values np ndarray tf tensor list tf tensor dict str tf tensor or dict str np ndarray and each example must have the shape batch_size num_channels height width pixel values pixel values can be obtained using autoimageprocessor see vitimageprocessor __call__ for details head_mask np ndarray or tf tensor of shape num_heads or num_layers num_heads optional mask to nullify selected heads of the self attention modules mask values selected in 0 1 1 indicates the head is not masked 0 indicates the head is masked output_attentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead output_hidden_states bool optional whether or not to return the hidden states of all layers see hidden_states under returned tensors for more detail this argument can be used only in eager mode in graph mode the value in the config will be used instead return_dict bool optional whether or not to return a file_utils modeloutput instead of a plain tuple this argument can be used in eager mode in graph mode the value will always be set to true training bool optional defaults to false whether or not to use the model in training mode some modules like dropout modules have different behaviors between training and evaluation returns examples python from transformers import autoimageprocessor tfvitmaemodel from pil import image import requests url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw image_processor autoimageprocessor from_pretrained facebook vit mae base model tfvitmaemodel from_pretrained facebook vit mae base inputs image_processor images image return_tensors tf outputs model inputs last_hidden_states outputs last_hidden_state encoder to decoder embed tokens append mask tokens to sequence no cls token unshuffle append cls token add pos embed apply transformer layers blocks predictor projection remove cls token args pixel_values tf tensor of shape batch_size height width num_channels or batch_size num_channels height width pixel values returns tf tensor of shape batch_size num_patches patch_size 2 num_channels patchified pixel values make sure channels are last sanity checks patchify args patchified_pixel_values tf tensor of shape batch_size num_patches patch_size 2 num_channels patchified pixel values returns tf tensor of shape batch_size height width num_channels pixel values sanity check unpatchify args pixel_values tf tensor of shape batch_size height width num_channels pixel values pred tf tensor of shape batch_size num_patches patch_size 2 num_channels predicted pixel values mask tf tensor of shape batch_size sequence_length tensor indicating which patches are masked 1 and which are not 0 returns tf tensor pixel reconstruction loss batch_size num_patches mean loss per patch mean loss on removed patches returns examples python from transformers import autoimageprocessor tfvitmaeforpretraining from pil import image import requests url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw image_processor autoimageprocessor from_pretrained facebook vit mae base model tfvitmaeforpretraining from_pretrained facebook vit mae base inputs image_processor images image return_tensors pt outputs model inputs loss outputs loss mask outputs mask ids_restore outputs ids_restore batch_size num_patches patch_size 2 3
from __future__ import annotations

import collections.abc
import math
from copy import deepcopy
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import tensorflow as tf

from ...activations_tf import get_tf_activation
from ...file_utils import (
    ModelOutput,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    replace_return_docstrings,
)
from ...modeling_tf_outputs import TFBaseModelOutput
from ...modeling_tf_utils import (
    TFModelInputType,
    TFPreTrainedModel,
    get_initializer,
    keras_serializable,
    unpack_inputs,
)
from ...tf_utils import shape_list, stable_softmax
from ...utils import logging
from .configuration_vit_mae import ViTMAEConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "ViTMAEConfig"
_CHECKPOINT_FOR_DOC = "facebook/vit-mae-base"


@dataclass
class TFViTMAEModelOutput(ModelOutput):
    

    last_hidden_state: tf.Tensor = None
    mask: tf.Tensor = None
    ids_restore: tf.Tensor = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None


@dataclass
class TFViTMAEDecoderOutput(ModelOutput):
    

    logits: tf.Tensor = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None


@dataclass
class TFViTMAEForPreTrainingOutput(ModelOutput):
    

    loss: tf.Tensor | None = None
    logits: tf.Tensor = None
    mask: tf.Tensor = None
    ids_restore: tf.Tensor = None
    hidden_states: Tuple[tf.Tensor] | None = None
    attentions: Tuple[tf.Tensor] | None = None


def get_2d_sincos_pos_embed(embed_dim, grid_size, add_cls_token=False):
    
    grid_h = tf.range(grid_size, dtype=tf.float32)
    grid_w = tf.range(grid_size, dtype=tf.float32)
    grid = tf.meshgrid(grid_w, grid_h)  
    grid = tf.stack(grid, axis=0)

    grid = tf.reshape(grid, [2, 1, grid_size, grid_size])
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if add_cls_token:
        pos_embed = tf.concat([tf.zeros((1, embed_dim)), pos_embed], axis=0)
    return pos_embed


def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
    if embed_dim % 2 != 0:
        raise ValueError("embed_dim must be even")

    
    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  
    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  

    emb = tf.concat([emb_h, emb_w], axis=1)  
    return emb


def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
    
    if embed_dim % 2 != 0:
        raise ValueError("embed_dim must be even")

    omega = tf.range(embed_dim // 2, dtype="float32")
    omega /= embed_dim / 2.0
    omega = 1.0 / 10000**omega  

    pos = tf.reshape(pos, [-1])  
    out = tf.einsum("m,d->md", pos, omega)  

    
    
    emb_sin = tf.sin(out)  
    emb_cos = tf.cos(out)  

    emb = tf.concat([emb_sin, emb_cos], axis=1)  
    return emb


class TFViTMAEEmbeddings(tf.keras.layers.Layer):
    

    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.patch_embeddings = TFViTMAEPatchEmbeddings(config, name="patch_embeddings")
        self.num_patches = self.patch_embeddings.num_patches

        self.config = config

    def build(self, input_shape: tf.TensorShape):
        self.cls_token = self.add_weight(
            shape=(1, 1, self.config.hidden_size),
            initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),
            trainable=True,
            name="cls_token",
        )
        self.position_embeddings = self.add_weight(
            shape=(1, self.num_patches + 1, self.config.hidden_size),
            initializer="zeros",
            trainable=False,  
            name="position_embeddings",
        )
        pos_embed = get_2d_sincos_pos_embed(
            self.position_embeddings.shape[-1],
            int(self.patch_embeddings.num_patches**0.5),
            add_cls_token=True,
        )[None, ...]
        self.position_embeddings.assign(pos_embed)

        super().build(input_shape)

    def random_masking(self, sequence: tf.Tensor, noise: tf.Tensor | None = None):
        
        batch_size, seq_length, dim = shape_list(sequence)
        len_keep = int(seq_length * (1 - self.config.mask_ratio))

        if noise is None:
            noise = tf.random.uniform(shape=(batch_size, seq_length), minval=0.0, maxval=1.0)  

        
        ids_shuffle = tf.argsort(noise, axis=1)  
        ids_restore = tf.argsort(ids_shuffle, axis=1)

        
        ids_keep = ids_shuffle[:, :len_keep]
        sequence_unmasked = tf.gather(
            sequence,
            axis=1,
            batch_dims=1,
            indices=ids_keep,
        )

        
        
        
        mask_keep = tf.zeros((batch_size, len_keep))
        mask_remove = tf.ones((batch_size, seq_length - len_keep))
        mask = tf.concat([mask_keep, mask_remove], axis=-1)

        
        mask = tf.gather(mask, axis=1, batch_dims=1, indices=ids_restore)

        return sequence_unmasked, mask, ids_restore

    def call(self, pixel_values: tf.Tensor, noise: tf.Tensor = None) -> tf.Tensor:
        embeddings = self.patch_embeddings(pixel_values)

        
        embeddings = embeddings + self.position_embeddings[:, 1:, :]

        
        embeddings, mask, ids_restore = self.random_masking(embeddings, noise)

        
        cls_token = self.cls_token + self.position_embeddings[:, :1, :]
        cls_tokens = tf.tile(cls_token, (shape_list(embeddings)[0], 1, 1))
        embeddings = tf.concat([cls_tokens, embeddings], axis=1)

        return embeddings, mask, ids_restore


class TFViTMAEPatchEmbeddings(tf.keras.layers.Layer):
    

    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)
        image_size, patch_size = config.image_size, config.patch_size
        num_channels, hidden_size = config.num_channels, config.hidden_size
        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.num_channels = num_channels
        self.config = config

        self.projection = tf.keras.layers.Conv2D(
            filters=hidden_size,
            kernel_size=patch_size,
            strides=patch_size,
            padding="valid",
            data_format="channels_last",
            kernel_initializer="glorot_uniform",  
            bias_initializer="zeros",
            name="projection",
        )

    def call(self, pixel_values: tf.Tensor, training: bool = False) -> tf.Tensor:
        batch_size, num_channels, height, width = shape_list(pixel_values)
        if tf.executing_eagerly():
            if num_channels != self.num_channels:
                raise ValueError(
                    "Make sure that the channel dimension of the pixel values match with the one set in the"
                    " configuration."
                )
            if height != self.image_size[0] or width != self.image_size[1]:
                raise ValueError(
                    f"Input image size ({height}*{width}) doesn't match model"
                    f" ({self.image_size[0]}*{self.image_size[1]})."
                )

        
        
        
        pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))

        projection = self.projection(pixel_values)

        
        
        num_patches = (width // self.patch_size[1]) * (height // self.patch_size[0])
        x = tf.reshape(tensor=projection, shape=(batch_size, num_patches, -1))

        return x



class TFViTMAESelfAttention(tf.keras.layers.Layer):
    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number "
                f"of attention heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size
        self.sqrt_att_head_size = math.sqrt(self.attention_head_size)

        self.query = tf.keras.layers.Dense(
            units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"
        )
        self.key = tf.keras.layers.Dense(
            units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"
        )
        self.value = tf.keras.layers.Dense(
            units=self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"
        )
        self.dropout = tf.keras.layers.Dropout(rate=config.attention_probs_dropout_prob)

    def transpose_for_scores(self, tensor: tf.Tensor, batch_size: int) -> tf.Tensor:
        
        tensor = tf.reshape(tensor=tensor, shape=(batch_size, -1, self.num_attention_heads, self.attention_head_size))

        
        return tf.transpose(tensor, perm=[0, 2, 1, 3])

    def call(
        self,
        hidden_states: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(inputs=hidden_states)
        mixed_key_layer = self.key(inputs=hidden_states)
        mixed_value_layer = self.value(inputs=hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        
        
        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(self.sqrt_att_head_size, dtype=attention_scores.dtype)
        attention_scores = tf.divide(attention_scores, dk)

        
        attention_probs = stable_softmax(logits=attention_scores, axis=-1)

        
        
        attention_probs = self.dropout(inputs=attention_probs, training=training)

        
        if head_mask is not None:
            attention_probs = tf.multiply(attention_probs, head_mask)

        attention_output = tf.matmul(attention_probs, value_layer)
        attention_output = tf.transpose(attention_output, perm=[0, 2, 1, 3])

        
        attention_output = tf.reshape(tensor=attention_output, shape=(batch_size, -1, self.all_head_size))
        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)

        return outputs



class TFViTMAESelfOutput(tf.keras.layers.Layer):
    

    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.dense = tf.keras.layers.Dense(
            units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
        )
        self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)

    def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
        hidden_states = self.dense(inputs=hidden_states)
        hidden_states = self.dropout(inputs=hidden_states, training=training)

        return hidden_states



class TFViTMAEAttention(tf.keras.layers.Layer):
    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.self_attention = TFViTMAESelfAttention(config, name="attention")
        self.dense_output = TFViTMAESelfOutput(config, name="output")

    def prune_heads(self, heads):
        raise NotImplementedError

    def call(
        self,
        input_tensor: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        self_outputs = self.self_attention(
            hidden_states=input_tensor, head_mask=head_mask, output_attentions=output_attentions, training=training
        )
        attention_output = self.dense_output(
            hidden_states=self_outputs[0], input_tensor=input_tensor, training=training
        )
        outputs = (attention_output,) + self_outputs[1:]  

        return outputs



class TFViTMAEIntermediate(tf.keras.layers.Layer):
    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.dense = tf.keras.layers.Dense(
            units=config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
        )

        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = get_tf_activation(config.hidden_act)
        else:
            self.intermediate_act_fn = config.hidden_act

    def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
        hidden_states = self.dense(inputs=hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)

        return hidden_states



class TFViTMAEOutput(tf.keras.layers.Layer):
    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.dense = tf.keras.layers.Dense(
            units=config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
        )
        self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)

    def call(self, hidden_states: tf.Tensor, input_tensor: tf.Tensor, training: bool = False) -> tf.Tensor:
        hidden_states = self.dense(inputs=hidden_states)
        hidden_states = self.dropout(inputs=hidden_states, training=training)
        hidden_states = hidden_states + input_tensor

        return hidden_states



class TFViTMAELayer(tf.keras.layers.Layer):
    

    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.attention = TFViTMAEAttention(config, name="attention")
        self.intermediate = TFViTMAEIntermediate(config, name="intermediate")
        self.vit_output = TFViTMAEOutput(config, name="output")

        self.layernorm_before = tf.keras.layers.LayerNormalization(
            epsilon=config.layer_norm_eps, name="layernorm_before"
        )
        self.layernorm_after = tf.keras.layers.LayerNormalization(
            epsilon=config.layer_norm_eps, name="layernorm_after"
        )

    def call(
        self,
        hidden_states: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        training: bool = False,
    ) -> Tuple[tf.Tensor]:
        attention_outputs = self.attention(
            
            input_tensor=self.layernorm_before(inputs=hidden_states),
            head_mask=head_mask,
            output_attentions=output_attentions,
            training=training,
        )
        attention_output = attention_outputs[0]

        
        hidden_states = attention_output + hidden_states

        
        layer_output = self.layernorm_after(inputs=hidden_states)

        intermediate_output = self.intermediate(hidden_states=layer_output)

        
        layer_output = self.vit_output(
            hidden_states=intermediate_output, input_tensor=hidden_states, training=training
        )
        outputs = (layer_output,) + attention_outputs[1:]  

        return outputs



class TFViTMAEEncoder(tf.keras.layers.Layer):
    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.layer = [TFViTMAELayer(config, name=f"layer_._{i}") for i in range(config.num_hidden_layers)]

    def call(
        self,
        hidden_states: tf.Tensor,
        head_mask: tf.Tensor,
        output_attentions: bool,
        output_hidden_states: bool,
        return_dict: bool,
        training: bool = False,
    ) -> Union[TFBaseModelOutput, Tuple[tf.Tensor]]:
        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(
                hidden_states=hidden_states,
                head_mask=head_mask[i],
                output_attentions=output_attentions,
                training=training,
            )
            hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

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

        return TFBaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
        )


@keras_serializable
class TFViTMAEMainLayer(tf.keras.layers.Layer):
    config_class = ViTMAEConfig

    def __init__(self, config: ViTMAEConfig, **kwargs):
        super().__init__(**kwargs)

        self.config = config

        self.embeddings = TFViTMAEEmbeddings(config, name="embeddings")
        self.encoder = TFViTMAEEncoder(config, name="encoder")
        self.layernorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layernorm")

    def get_input_embeddings(self) -> tf.keras.layers.Layer:
        return self.embeddings.patch_embeddings

    def _prune_heads(self, heads_to_prune):
        
        raise NotImplementedError

    @unpack_inputs
    def call(
        self,
        pixel_values: TFModelInputType | None = None,
        noise: tf.Tensor = None,
        head_mask: np.ndarray | tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFViTMAEModelOutput, Tuple[tf.Tensor]]:
        embedding_output, mask, ids_restore = self.embeddings(
            pixel_values=pixel_values, training=training, noise=noise
        )

        
        
        
        
        
        if head_mask is not None:
            raise NotImplementedError
        else:
            head_mask = [None] * self.config.num_hidden_layers

        encoder_outputs = self.encoder(
            embedding_output,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        sequence_output = encoder_outputs[0]
        sequence_output = self.layernorm(inputs=sequence_output)

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

        return TFViTMAEModelOutput(
            last_hidden_state=sequence_output,
            mask=mask,
            ids_restore=ids_restore,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )


class TFViTMAEPreTrainedModel(TFPreTrainedModel):
    

    config_class = ViTMAEConfig
    base_model_prefix = "vit"
    main_input_name = "pixel_values"


VIT_MAE_START_DOCSTRING = r

VIT_MAE_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "The bare ViTMAE Model transformer outputting raw hidden-states without any specific head on top.",
    VIT_MAE_START_DOCSTRING,
)
class TFViTMAEModel(TFViTMAEPreTrainedModel):
    def __init__(self, config: ViTMAEConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

        self.vit = TFViTMAEMainLayer(config, name="vit")

    def get_input_embeddings(self):
        return self.vit.get_input_embeddings()

    @unpack_inputs
    @add_start_docstrings_to_model_forward(VIT_MAE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFViTMAEModelOutput, config_class=_CONFIG_FOR_DOC)
    def call(
        self,
        pixel_values: TFModelInputType | None = None,
        noise: tf.Tensor = None,
        head_mask: np.ndarray | tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFViTMAEModelOutput, Tuple[tf.Tensor]]:
        r
        outputs = self.vit(
            pixel_values=pixel_values,
            noise=noise,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        return outputs


class TFViTMAEDecoder(tf.keras.layers.Layer):
    def __init__(self, config, num_patches, **kwargs):
        super().__init__(**kwargs)
        self.decoder_embed = tf.keras.layers.Dense(config.decoder_hidden_size, name="decoder_embed")

        decoder_config = deepcopy(config)
        decoder_config.hidden_size = config.decoder_hidden_size
        decoder_config.num_hidden_layers = config.decoder_num_hidden_layers
        decoder_config.num_attention_heads = config.decoder_num_attention_heads
        decoder_config.intermediate_size = config.decoder_intermediate_size
        self.decoder_layers = [
            TFViTMAELayer(decoder_config, name=f"decoder_layers.{j}") for j in range(config.decoder_num_hidden_layers)
        ]

        self.decoder_norm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="decoder_norm")
        self.decoder_pred = tf.keras.layers.Dense(
            config.patch_size**2 * config.num_channels,
            kernel_initializer=get_initializer(config.initializer_range),
            name="decoder_pred",
        )  
        self.config = config
        self.num_patches = num_patches

    def build(self, input_shape: tf.TensorShape):
        self.mask_token = self.add_weight(
            shape=(1, 1, self.config.decoder_hidden_size),
            initializer=tf.random_normal_initializer(stddev=self.config.initializer_range),
            trainable=True,
            name="mask_token",
        )
        self.decoder_pos_embed = self.add_weight(
            shape=(1, self.num_patches + 1, self.config.decoder_hidden_size),
            initializer="zeros",
            trainable=False,
            name="decoder_pos_embed",
        )
        decoder_pos_embed = get_2d_sincos_pos_embed(
            self.decoder_pos_embed.shape[-1],
            int(self.num_patches**0.5),
            add_cls_token=True,
        )[None, ...]
        self.decoder_pos_embed.assign(decoder_pos_embed)

        super().build(input_shape)

    def call(
        self,
        hidden_states,
        ids_restore,
        output_attentions=False,
        output_hidden_states=False,
        return_dict=True,
    ):
        
        x = self.decoder_embed(hidden_states)

        
        mask_tokens = tf.tile(
            self.mask_token,
            (shape_list(x)[0], shape_list(ids_restore)[1] + 1 - shape_list(x)[1], 1),
        )
        x_ = tf.concat([x[:, 1:, :], mask_tokens], axis=1)  
        x_ = tf.gather(x_, axis=1, batch_dims=1, indices=ids_restore)  
        x = tf.concat([x[:, :1, :], x_], axis=1)  

        
        hidden_states = x + self.decoder_pos_embed

        
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        for i, layer_module in enumerate(self.decoder_layers):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(
                hidden_states,
                head_mask=None,
                output_attentions=output_attentions,
            )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        hidden_states = self.decoder_norm(hidden_states)

        
        logits = self.decoder_pred(hidden_states)

        
        logits = logits[:, 1:, :]

        if not return_dict:
            return tuple(v for v in [logits, all_hidden_states, all_self_attentions] if v is not None)
        return TFViTMAEDecoderOutput(logits=logits, hidden_states=all_hidden_states, attentions=all_self_attentions)


@add_start_docstrings(
    "The ViTMAE Model transformer with the decoder on top for self-supervised pre-training.",
    VIT_MAE_START_DOCSTRING,
)
class TFViTMAEForPreTraining(TFViTMAEPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config

        self.vit = TFViTMAEMainLayer(config, name="vit")
        self.decoder = TFViTMAEDecoder(
            config,
            num_patches=self.vit.embeddings.num_patches,
            name="decoder",
        )

    def get_input_embeddings(self):
        return self.vit.get_input_embeddings()

    def _prune_heads(self, heads_to_prune):
        raise NotImplementedError

    def patchify(self, pixel_values):
        
        patch_size, num_channels = self.config.patch_size, self.config.num_channels
        
        if shape_list(pixel_values)[1] == num_channels:
            pixel_values = tf.transpose(pixel_values, perm=(0, 2, 3, 1))

        
        tf.debugging.assert_equal(
            shape_list(pixel_values)[1],
            shape_list(pixel_values)[2],
            message="Make sure the pixel values have a squared size",
        )
        tf.debugging.assert_equal(
            shape_list(pixel_values)[1] % patch_size,
            0,
            message="Make sure the pixel values have a size that is divisible by the patch size",
        )
        tf.debugging.assert_equal(
            shape_list(pixel_values)[3],
            num_channels,
            message=(
                "Make sure the number of channels of the pixel values is equal to the one set in the configuration"
            ),
        )

        
        batch_size = shape_list(pixel_values)[0]
        num_patches_one_direction = shape_list(pixel_values)[2] // patch_size
        patchified_pixel_values = tf.reshape(
            pixel_values,
            (batch_size, num_patches_one_direction, patch_size, num_patches_one_direction, patch_size, num_channels),
        )
        patchified_pixel_values = tf.einsum("nhpwqc->nhwpqc", patchified_pixel_values)
        patchified_pixel_values = tf.reshape(
            patchified_pixel_values,
            (batch_size, num_patches_one_direction * num_patches_one_direction, patch_size**2 * num_channels),
        )
        return patchified_pixel_values

    def unpatchify(self, patchified_pixel_values):
        
        patch_size, num_channels = self.config.patch_size, self.config.num_channels
        num_patches_one_direction = int(shape_list(patchified_pixel_values)[1] ** 0.5)
        
        tf.debugging.assert_equal(
            num_patches_one_direction * num_patches_one_direction,
            shape_list(patchified_pixel_values)[1],
            message="Make sure that the number of patches can be squared",
        )

        
        batch_size = shape_list(patchified_pixel_values)[0]
        patchified_pixel_values = tf.reshape(
            patchified_pixel_values,
            (batch_size, num_patches_one_direction, num_patches_one_direction, patch_size, patch_size, num_channels),
        )
        patchified_pixel_values = tf.einsum("nhwpqc->nhpwqc", patchified_pixel_values)
        pixel_values = tf.reshape(
            patchified_pixel_values,
            (batch_size, num_patches_one_direction * patch_size, num_patches_one_direction * patch_size, num_channels),
        )
        return pixel_values

    def forward_loss(self, pixel_values, pred, mask):
        
        target = self.patchify(pixel_values)
        if self.config.norm_pix_loss:
            mean = tf.reduce_mean(target, axis=-1, keepdims=True)
            var = tf.math.reduce_variance(target, axis=-1, keepdims=True)
            target = (target - mean) / (var + 1.0e-6) ** 0.5

        loss = (pred - target) ** 2
        loss = tf.reduce_mean(loss, axis=-1)  

        loss = tf.reduce_sum(loss * mask) / tf.reduce_sum(mask)  
        loss = tf.reshape(loss, (1,))
        return loss

    @unpack_inputs
    @add_start_docstrings_to_model_forward(VIT_MAE_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=TFViTMAEForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
    def call(
        self,
        pixel_values: TFModelInputType | None = None,
        noise: tf.Tensor = None,
        head_mask: np.ndarray | tf.Tensor | None = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        training: bool = False,
    ) -> Union[TFViTMAEForPreTrainingOutput, Tuple[tf.Tensor]]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.vit(
            pixel_values=pixel_values,
            noise=noise,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            training=training,
        )

        latent = outputs.last_hidden_state
        ids_restore = outputs.ids_restore
        mask = outputs.mask

        decoder_outputs = self.decoder(latent, ids_restore)  
        logits = decoder_outputs.logits

        loss = self.forward_loss(pixel_values, logits, mask)

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

        return TFViTMAEForPreTrainingOutput(
            loss=loss,
            logits=logits,
            mask=mask,
            ids_restore=ids_restore,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )