text: import math
  from typing import List, Union
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch import Tensor
  
  
  @torch.jit.script
  def spatial_pyramid_pool(
      input: Tensor, bins: Union[int, List[int]], mode: str = "max"
  ) -> Tensor:
      """Spatial Pyramid Pooling: https://arxiv.org/pdf/1406.4729.pdf
  
      Args:
          input (Tensor): an input tensor expected from the convolutional layer.
          bins (List[int]): a list of integer of preferred size of outputs.
          mode (str): how to reduce the spatial space.
  
      Returns:
          outputs (Tensor): a flatten tensor with size (batch, bins[0] * bins[0] + bins[1]
          * bins[1] + ...)
      """
  
      assert mode in ["max", "mean", "average", "avg"]
      b, _, h, w = input.shape
      bins = [bins] if isinstance(bins, int) else bins
      outputs = []
  
      for bin_ in bins:
          h_kernel = math.ceil(h / bin_)
          w_kernel = math.ceil(w / bin_)
          h_stride = math.floor(h / bin_)
          w_stride = math.floor(w / bin_)
  
          if mode == "max":
              output = F.max_pool2d(
                  input, kernel_size=(h_kernel, w_kernel), stride=(h_stride, w_stride)
              )
          else:
              output = F.avg_pool2d(
                  input, kernel_size=(h_kernel, w_kernel), stride=(h_stride, w_stride)
              )
          output = output.view(b, -1)
          outputs.append(output)
  
      outputs = torch.cat(outputs, dim=-1)
      return outputs
  
  
  class SpaitalPyramidPool(nn.Module):
      def __init__(self, bins: Union[int, List[int]], mode: str = "max") -> None:
          super().__init__()
          self.bins = bins
          self.mode = mode
  
      def forward(self, input: Tensor) -> Tensor:
          return spatial_pyramid_pool(input, bins=self.bins, mode=self.mode)
  
  
  if __name__ == "__main__":
      input = torch.zeros(1, 512, 13, 13)
      output = spatial_pyramid_pool(input, [1, 2, 3], "max")
      print(output.shape)
  