text: import operator_benchmark as op_bench
  import torch
  
  """Microbenchmarks for interpolate operator."""
  
  
  class InterpolateBenchmark(op_bench.TorchBenchmarkBase):
      def init(self, input_size, output_size, channels_last=False, mode='linear', dtype=torch.float):
  
          input_image = torch.randint(0, 256, size=input_size, dtype=dtype, device='cpu',
                                      requires_grad=self.auto_set())
          if channels_last:
              if input_image.ndim == 4:
                  input_image = input_image.contiguous(memory_format=torch.channels_last)
              elif input_image.ndim == 5:
                  input_image = input_image.contiguous(memory_format=torch.channels_last_3d)
              else:
                  raise ValueError(
                      f"Can not set channels_last to the input of {input_image.ndim} dims"
                  )
  
  
          align_corners = None if mode == "nearest" else False
  
          if mode == "linear":
              mode = {
                  3: 'linear',
                  4: 'bilinear',
                  5: 'trilinear',
              }[input_image.ndim]
  
          self.inputs = {
              "input_image": input_image,
              "output_size": output_size,
              "mode": mode,
              "align_corners": align_corners,
          }
  
          self.set_module_name("interpolate")
  
      def forward(self, input_image, output_size, mode, align_corners):
          return torch.nn.functional.interpolate(input_image, size=output_size, mode=mode,
                                                 align_corners=align_corners)
  
  
  config_short = op_bench.config_list(
      attr_names=["input_size", "output_size"],
      attrs=[
          [(1, 3, 60, 40), (24, 24)],
          [(1, 3, 600, 400), (240, 240)],
          [(1, 3, 320, 320), (256, 256)],
      ],
      cross_product_configs={
          'channels_last': [True, False],
          'mode': ["nearest", "linear", "bicubic"],
      },
      tags=["short"],
  )
  
  config_short += op_bench.config_list(
      attr_names=["input_size", "output_size"],
      attrs=[
          [(1, 3, 60, 40), (24, 24)],
          [(1, 3, 600, 400), (240, 240)],
          [(1, 3, 320, 320), (256, 256)],
      ],
      cross_product_configs={
          'channels_last': [True, False],
          'mode': ["nearest", ],
          'dtype': [torch.uint8, ],
      },
      tags=["short"],
  )
  
  
  config_long = op_bench.config_list(
      attr_names=["input_size", "output_size"],
      attrs=[
          [(1, 3, 320, 320), (512, 512)],
          [(1, 3, 500, 500), (256, 256)],
          [(1, 3, 500, 500), (800, 800)],
  
          # vectorization test-case
          [(2, 128, 64, 46), (128, 128)],
          [(2, 128, 64, 46), (32, 24)],
      ],
      cross_product_configs={
          'channels_last': [True, False],
          'mode': ["nearest", "linear", "bicubic"],
      },
      tags=["long"],
  )
  
  
  config_3d = op_bench.config_list(
      # no channels_last for 3D tensors
      attr_names=["input_size", "output_size"],
      attrs=[
          [(4, 512, 320), (256,)],
          [(4, 512, 320), (512,)],
      ],
      cross_product_configs={
          'mode': ["nearest", "linear"],
      },
      tags=["long"],
  )
  
  
  config_5d = op_bench.config_list(
      attr_names=["input_size", "output_size"],
      attrs=[
          [(1, 3, 16, 320, 320), (8, 256, 256)],
          [(1, 3, 16, 320, 320), (32, 512, 512)],
  
          # vectorization test-case
          [(1, 16, 32, 64, 64), (16, 32, 32)],
          [(1, 16, 32, 64, 64), (64, 128, 128)],
      ],
      cross_product_configs={
          'channels_last': [True, False],
          'mode': ["nearest", "linear"],
      },
      tags=["long"],
  )
  
  
  for config in (config_short, config_long, config_3d, config_5d):
      op_bench.generate_pt_test(config, InterpolateBenchmark)
  
  
  if __name__ == "__main__":
      op_bench.benchmark_runner.main()
  