code: def test_equivalence(art_warning, fix_get_cifar10_data):
      import torch
      from art.estimators.certification.derandomized_smoothing import (
          PyTorchDeRandomizedSmoothing,
      )
      from art.estimators.certification.derandomized_smoothing.vision_transformers.pytorch import (
          PyTorchVisionTransformer,
      )
  
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
  
      try:
  
          class MadrylabImplementations:
              """
              Code adapted from the implementation in https://github.com/MadryLab/smoothed-vit
              to check against our own functionality.
  
              Original License:
  
              MIT License
  
              Copyright (c) 2021 Madry Lab
  
              Permission is hereby granted, free of charge, to any person obtaining a copy
              of this software and associated documentation files (the "Software"), to deal
              in the Software without restriction, including without limitation the rights
              to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
              copies of the Software, and to permit persons to whom the Software is
              furnished to do so, subject to the following conditions:
  
              The above copyright notice and this permission notice shall be included in all
              copies or substantial portions of the Software.
  
              THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
              IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
              FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
              AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
              LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
              OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
              SOFTWARE.
  
              """
  
              def __init__(self):
                  pass
  
              @classmethod
              def token_dropper(cls, x, mask):
                  """
                  The implementation of dropping tokens has been done slightly differently in this tool.
                  Here we check that it is equivalent to the original implementation
                  """
  
                  class MaskProcessor(torch.nn.Module):
                      def __init__(self, patch_size=16):
                          super().__init__()
                          self.avg_pool = torch.nn.AvgPool2d(patch_size)
  
                      def forward(self, ones_mask):
                          B = ones_mask.shape[0]
                          ones_mask = ones_mask[0].unsqueeze(
                              0
                          )  # take the first mask
                          ones_mask = self.avg_pool(ones_mask)[0]
                          ones_mask = (
                              torch.where(ones_mask.view(-1) > 0)[0] + 1
                          )
                          ones_mask = torch.cat(
                              [
                                  torch.IntTensor(1)
                                  .fill_(0)
                                  .to(device),
                                  ones_mask,
                              ]
                          ).unsqueeze(0)
                          ones_mask = ones_mask.expand(B, -1)
                          return ones_mask
  
                  mask_processor = MaskProcessor()
                  patch_mask = mask_processor(mask)
  
                  # x = self.pos_drop(x) # B, N, C
                  if patch_mask is not None:
                      # patch_mask is B, K
                      B, N, C = x.shape
                      if (
                          len(patch_mask.shape) == 1
                      ):  # not a separate one per batch
                          x = x[:, patch_mask]
                      else:
                          patch_mask = patch_mask.unsqueeze(-1).expand(
                              -1, -1, C
                          )
                          x = torch.gather(x, 1, patch_mask)
                  return x
  
              @classmethod
              def embedder(cls, x, pos_embed, cls_token):
                  """
                  NB, original code used the pos embed from the divit rather than vit
                  (which we pull from our model) which we use here.
  
                  From timm vit:
                  self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * .02)
  
                  From timm dvit:
                  self.pos_embed = nn.Parameter(torch.zeros(1,
                                                           self.patch_embed.num_patches + self.num_prefix_tokens,
                                                           self.embed_dim))
  
                  From repo:
                  self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
                  """
                  x = torch.cat(
                      (cls_token.expand(x.shape[0], -1, -1), x), dim=1
                  )
                  return x + pos_embed
  
          def forward_features(self, x: torch.Tensor) -> torch.Tensor:
              """
              This is a copy of the function in ArtViT.forward_features
              except we also perform an equivalence assertion compared to the implementation
              in https://github.com/MadryLab/smoothed-vit (see MadrylabImplementations class above)
  
              The forward pass of the ViT.
  
              :param x: Input data.
              :return: The input processed by the ViT backbone
              """
              import copy
  
              ablated_input = False
              if x.shape[1] == self.in_chans + 1:
                  ablated_input = True
  
              if ablated_input:
                  x, ablation_mask = (
                      x[:, : self.in_chans],
                      x[:, self.in_chans : self.in_chans + 1],
                  )
  
              x = self.patch_embed(x)
  
              madry_embed = MadrylabImplementations.embedder(
                  copy.copy(x), self.pos_embed, self.cls_token
              )
              x = self._pos_embed(x)
              assert torch.equal(madry_embed, x)
  
              # pass the x into the token dropping code
              madry_dropped = MadrylabImplementations.token_dropper(
                  copy.copy(x), ablation_mask
              )
  
              if self.to_drop_tokens and ablated_input:
                  ones = self.ablation_mask_embedder(ablation_mask)
                  to_drop = torch.sum(ones, dim=2)
                  indexes = torch.gt(torch.where(to_drop > 1, 1, 0), 0)
                  x = self.drop_tokens(x, indexes)
  
              assert torch.equal(madry_dropped, x)
  
              x = self.norm_pre(x)
              x = self.blocks(x)
  
              return self.norm(x)
  
          # Replace the forward_features with the forward_features code with checks.
          PyTorchVisionTransformer.forward_features = forward_features
  
          art_model = PyTorchDeRandomizedSmoothing(
              model="vit_small_patch16_224",
              loss=torch.nn.CrossEntropyLoss(),
              optimizer=torch.optim.SGD,
              optimizer_params={"lr": 0.01},
              input_shape=(3, 32, 32),
              nb_classes=10,
              ablation_size=4,
              load_pretrained=False,
              replace_last_layer=True,
              verbose=False,
          )
  
          cifar_data = fix_get_cifar10_data[0][:50]
          cifar_labels = fix_get_cifar10_data[1][:50]
  
          scheduler = torch.optim.lr_scheduler.MultiStepLR(
              art_model.optimizer, milestones=[1], gamma=0.1
          )
          art_model.fit(
              cifar_data,
              cifar_labels,
              nb_epochs=1,
              update_batchnorm=True,
              scheduler=scheduler,
              batch_size=128,
          )
      except ARTTestException as e:
          art_warning(e)
  
file_path: adversarial-robustness-toolbox/tests/estimators/certification/test_vision_transformers.py
imports: import pytest
  import os
  import numpy
  from art.utils import load_dataset
  from tests.utils import ARTTestException
label_desc: Use image node_id 9 for calling a global function with example usage: test_equivalence(art_warning, fix_get_cifar10_data) and returns: self, x, unknown, ones_mask