decomp_upsample_bilinear_aa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9baca3ca-596c-420c-bb44-0f8cc55e5dd4",
   "metadata": {},
   "source": [
    "## Try to write a decomposition for upsampling bilinear AA=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "54d52c86-638e-4f91-8e7b-b87067e324f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5ddc6cc9-5219-4e60-b91a-ba65ebba07ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "n, c, in_h, in_w = 5, 3, 234, 345\n",
    "out_h, out_w = 123, 234\n",
    "\n",
    "align_corners = False\n",
    "\n",
    "input_tensor = torch.arange(n * c * in_h * in_w, dtype=torch.float32).reshape(n, c, in_h, in_w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f881cd9b-e7e8-462c-aa73-304170dc57da",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6fa07b82-fd4c-4d21-884c-b111d8af8549",
   "metadata": {},
   "outputs": [],
   "source": [
    "# UpSample.h\n",
    "# template <typename scalar_t>\n",
    "# static inline scalar_t area_pixel_compute_scale(\n",
    "#     int64_t input_size,\n",
    "#     int64_t output_size,\n",
    "#     bool align_corners,\n",
    "#     const c10::optional<double> scale) {\n",
    "#   // see Note [area_pixel_compute_scale]\n",
    "#   if(align_corners) {\n",
    "#     if(output_size > 1) {\n",
    "#       return static_cast<scalar_t>(input_size - 1) / (output_size - 1);\n",
    "#     } else {\n",
    "#       return static_cast<scalar_t>(0);\n",
    "#     }\n",
    "#   } else {\n",
    "#     return compute_scales_value<scalar_t>(scale, input_size, output_size);\n",
    "#   }\n",
    "# }\n",
    "#\n",
    "# Same as compute_scale in decompositions.py::upsample_bicubic2d_default\n",
    "#\n",
    "def _area_pixel_compute_scale(in_size, out_size, align_corners, scale=None):\n",
    "    if align_corners:\n",
    "        return (in_size - 1) / (out_size - 1) if out_size > 1 else 0\n",
    "    else:\n",
    "        return 1 / scale if scale is not None and scale > 0 else in_size / out_size        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f05c12c-bea0-4502-8270-d2ca3eb2011a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "83dcf825-7d17-4c92-9201-8a0c8700f6a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def aa_linear_filter(x):\n",
    "    x = torch.abs(x)\n",
    "    return 1.0 - torch.clamp(x, max=1.0)\n",
    "\n",
    "\n",
    "def _compute_indices_weights_aa(out_size, in_size, scale, interp_size, device):\n",
    "    scale = _area_pixel_compute_scale(in_size, out_size, align_corners, scale=scale)\n",
    "        \n",
    "    support = torch.tensor((interp_size * 0.5) * scale if scale >= 1.0 else interp_size * 0.5, device=device)\n",
    "    max_interp_size = torch.ceil(support).to(torch.long) * 2 + 1\n",
    "\n",
    "    i = torch.arange(out_size, dtype=torch.long, device=device)\n",
    "\n",
    "    center = scale * (i + 0.5)\n",
    "    invscale = 1.0 / scale if scale >= 1.0 else 1.0\n",
    "\n",
    "    # compute source indices as [xmin, xmin+1, ..., xmin+xsize-1]\n",
    "    xmin = torch.clamp((center - support + 0.5).to(torch.long), min=0)\n",
    "    xsize = torch.clamp((center + support + 0.5).to(torch.long), max=in_size) - xmin\n",
    "    xsize = torch.clamp(xsize, 0, max_interp_size)\n",
    "    \n",
    "    # compute weights\n",
    "    j = torch.arange(max_interp_size, dtype=torch.long, device=device).view(-1, 1)\n",
    "    # TODO: use a generic function aa_filter defined for bilinear and bicubic\n",
    "    weights = aa_linear_filter((j + xmin - center + 0.5) * invscale)\n",
    "    weights = torch.where(j < xsize, weights, 0.0)\n",
    "    total_weights = weights.sum(dim=0)\n",
    "    weights = weights / total_weights\n",
    "\n",
    "    return xmin, xsize, weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "47306995-1565-4c73-b5e8-324d5f4aab2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial, reduce\n",
    "from typing import Callable, cast, Iterable, List, Optional, Tuple, Union\n",
    "from torch import sym_float, sym_int, Tensor\n",
    "\n",
    "\n",
    "def _sum_tensors(ts: Iterable[Tensor]) -> Tensor:\n",
    "    return reduce(torch.add, ts)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1d4583ab-1e69-4fc5-8067-0c6a70fad33f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _separable_upsample_bilinear2d_aa_single_dim(in_tensor, out_size, interp_dim, align_corners, scale=None):\n",
    "    # Assume that in_tensor dtype is float32\n",
    "    \n",
    "    assert interp_dim % 4 in (2, 3)\n",
    "    \n",
    "    n, c, in_h, in_w = in_tensor.shape\n",
    "    interp_size = 2  # bilinear\n",
    "    in_size = in_tensor.shape[interp_dim]\n",
    "        \n",
    "    n_idx = torch.arange(n, device=in_tensor.device).view(n, 1, 1, 1)\n",
    "    c_idx = torch.arange(c, device=in_tensor.device).view(1, c, 1, 1)\n",
    "    \n",
    "    if interp_dim % 4 == 3:\n",
    "        # horizontal pass\n",
    "        xmin, xsize, weights = _compute_indices_weights_aa(out_size, in_size, scale, interp_size, device=in_tensor.device)\n",
    "        in_y = torch.arange(in_h, device=in_tensor.device).view((1, 1, in_h, 1))\n",
    "        xmin_idx = xmin.view(1, 1, 1, out_size)\n",
    "        \n",
    "        max_interp_size = len(weights)\n",
    "        in_tensor_list = [in_tensor[n_idx, c_idx, in_y, torch.clamp(xmin_idx + k, max=in_w - 1)] for k in range(max_interp_size)]\n",
    "        w_tensor_list = weights.unbind(dim=0)\n",
    "        return _sum_tensors(in_t * w_t for in_t, w_t in zip(in_tensor_list, w_tensor_list))        \n",
    "    else:\n",
    "        # vertical pass\n",
    "        ymin, ysize, weights = _compute_indices_weights_aa(out_size, in_size, scale, interp_size, device=in_tensor.device)\n",
    "\n",
    "        ymin_idx = ymin.view(1, 1, out_size, 1)\n",
    "        in_x = torch.arange(in_w, device=in_tensor.device).view((1, 1, 1, in_w))\n",
    "\n",
    "        max_interp_size = len(weights)\n",
    "        in_tensor_list = [in_tensor[n_idx, c_idx, torch.clamp(ymin_idx + k, max=in_h - 1), in_x] for k in range(max_interp_size)]\n",
    "        w_tensor_list = weights.unsqueeze(-1).unbind(dim=0)\n",
    "        return _sum_tensors(in_t * w_t for in_t, w_t in zip(in_tensor_list, w_tensor_list))        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a10bfb6e-4432-45d5-a634-504398350d07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def upsample_bilinear2d(\n",
    "    input: Tensor,\n",
    "    output_size: List[int],\n",
    "    align_corners: bool,\n",
    "    scales_h: Optional[float] = None,\n",
    "    scales_w: Optional[float] = None,\n",
    ") -> Tensor:\n",
    "\n",
    "    # horizontal pass\n",
    "    if output_size[1] != input.shape[-1]:\n",
    "        output = _separable_upsample_bilinear2d_aa_single_dim(input, output_size[1], -1, align_corners=align_corners, scale=scales_w)\n",
    "    else:\n",
    "        output = input\n",
    "\n",
    "    # vertical pass\n",
    "    if output_size[0] != input.shape[-2]:\n",
    "        output = _separable_upsample_bilinear2d_aa_single_dim(output, output_size[0], -2, align_corners=align_corners, scale=scales_h)\n",
    "\n",
    "    return output    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "001921a8-ea8e-46e7-9d1b-5a10c540a0f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = upsample_bilinear2d(input_tensor, (out_h, out_w), align_corners=align_corners)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "84c4e567-12ab-4d1c-b3ac-e07e431e5cb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 3, 123, 234])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "efc8921d-68ed-4065-95cd-8dc20c0b6382",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([5, 3, 123, 234])\n",
      "tensor([[ 225.4809,  226.8456,  228.2820,  ...,  565.9496,  567.3861,\n",
      "          568.7507],\n",
      "        [ 806.9192,  808.2839,  809.7203,  ..., 1147.3879, 1148.8245,\n",
      "         1150.1890],\n",
      "        [1459.6218, 1460.9866, 1462.4229,  ..., 1800.0905, 1801.5269,\n",
      "         1802.8915],\n",
      "        [2112.3245, 2113.6890, 2115.1255,  ..., 2452.7932, 2454.2295,\n",
      "         2455.5938],\n",
      "        [2771.9033, 2773.2686, 2774.7046,  ..., 3112.3726, 3113.8088,\n",
      "         3115.1733]])\n",
      "tensor([[ 225.4810,  226.8455,  228.2820,  ...,  565.9496,  567.3861,\n",
      "          568.7506],\n",
      "        [ 806.9191,  808.2837,  809.7202,  ..., 1147.3878, 1148.8243,\n",
      "         1150.1890],\n",
      "        [1459.6216, 1460.9863, 1462.4229,  ..., 1800.0906, 1801.5266,\n",
      "         1802.8912],\n",
      "        [2112.3242, 2113.6890, 2115.1255,  ..., 2452.7935, 2454.2292,\n",
      "         2455.5938],\n",
      "        [2771.9033, 2773.2681, 2774.7048,  ..., 3112.3721, 3113.8088,\n",
      "         3115.1733]])\n"
     ]
    }
   ],
   "source": [
    "from torch.nn import functional as F\n",
    "\n",
    "expected = F.interpolate(input_tensor, size=(out_h, out_w), mode=\"bilinear\", align_corners=align_corners, antialias=True)\n",
    "\n",
    "print(expected.shape)\n",
    "print(expected[0, 0, :5, :])\n",
    "print(output[0, 0, :5, :])\n",
    "\n",
    "torch.testing.assert_close(expected, output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb25c16-ae6c-4a55-bcf3-219361059989",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33e6a1a3-5a41-425c-9216-43a248a23fd5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "30396228-bad2-4f74-bade-582b61b40159",
   "metadata": {},
   "source": [
    "### Horizontal pass development"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "02802c09-5e3a-49b3-9e93-d99d8e75ffad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.25"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scale = _area_pixel_compute_scale(in_w, out_w, align_corners)\n",
    "interp_size = 2  # bilinear\n",
    "\n",
    "scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "e685807f-f362-4547-8c07-6a67c617c5e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(6.2500), tensor(15))"
      ]
     },
     "execution_count": 156,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "support = torch.tensor((interp_size * 0.5) * scale if scale >= 1.0 else interp_size * 0.5, device=input_tensor.device)\n",
    "max_interp_size = torch.ceil(support).to(torch.long) * 2 + 1\n",
    "\n",
    "support, max_interp_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "56cbaa41-6de2-4f0c-885f-b25e5bc56ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = torch.arange(out_w, dtype=input_tensor.dtype, device=input_tensor.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "4edf04cf-d993-455c-9f5c-5156036a04e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 3.1250,  9.3750, 15.6250, 21.8750]), 0.16)"
      ]
     },
     "execution_count": 158,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "center = scale * (i + 0.5)\n",
    "invscale = 1.0 / scale if scale >= 1.0 else 1.0\n",
    "center, invscale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "id": "28628cad-180a-40f5-ac86-c147a50dd641",
   "metadata": {},
   "outputs": [],
   "source": [
    "# xmin = std::max(\n",
    "#         static_cast<int64_t>(center - support + 0.5 + align_corners_delta), static_cast<int64_t>(0));\n",
    "# xsize = std::min(\n",
    "#         static_cast<int64_t>(center + support + 0.5 + align_corners_delta), input_size) - xmin;\n",
    "\n",
    "xmin = torch.clamp((center - support + 0.5).to(torch.long), min=0)\n",
    "xsize = torch.clamp((center + support + 0.5).to(torch.long), max=in_w) - xmin\n",
    "xsize = torch.clamp(xsize, 0, max_interp_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "id": "93eb57ba-ee21-4b8f-a6f6-7cf87d5b0310",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 0,  3,  9, 16]), tensor([ 9, 13, 13,  9]))"
      ]
     },
     "execution_count": 185,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xmin, xsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "id": "27b8050a-ba8d-4e44-a29d-f192f3c0a8ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# template<typename scalar_t>\n",
    "# static inline scalar_t aa_filter(scalar_t x) {\n",
    "#   x = std::abs(x);\n",
    "#   if (x < 1.0) {\n",
    "#     return 1.0 - x;\n",
    "#   }\n",
    "#   return 0.0;\n",
    "# }\n",
    "\n",
    "# def aa_filter(x):\n",
    "#     x = torch.abs(x)\n",
    "#     return torch.where(x < 1, 1.0 - x, 0.0)\n",
    "\n",
    "def aa_filter(x):\n",
    "    x = torch.abs(x)\n",
    "    return 1.0 - torch.clamp(x, max=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "id": "c10c0a86-e21f-4865-afd6-ad77bfa85d5a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.0, 0.30000000000000004, 1.0, 0.30000000000000004, 0.0, 0.0)"
      ]
     },
     "execution_count": 187,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def aa_filter_scalar(x):\n",
    "    x = abs(x)\n",
    "    if x < 1.0:\n",
    "        return 1.0 - x\n",
    "    return 0.0\n",
    "\n",
    "\n",
    "aa_filter_scalar(-1.1), aa_filter_scalar(-0.7), aa_filter_scalar(0.0), aa_filter_scalar(0.7), aa_filter_scalar(1.1), aa_filter_scalar(1.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "id": "96a1d1cc-110e-4813-9518-fd0fbe727611",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0000, 0.3000, 1.0000, 0.3000, 0.0000, 0.0000])"
      ]
     },
     "execution_count": 188,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.tensor([-1.1, -0.7, 0.0, 0.7, 1.1, 1.0])\n",
    "\n",
    "aa_filter(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "id": "2180dabb-6d20-4329-a45c-4ea1e8ff6f92",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 9, 13, 13,  9])"
      ]
     },
     "execution_count": 189,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "id": "c065ccf2-b664-43fd-9838-12fb8f2aeb1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "j = torch.arange(max_interp_size, dtype=input_tensor.dtype, device=input_tensor.device).view(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "id": "8f213658-85e8-44b7-b146-f9df59b68874",
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = aa_filter((j + xmin - center + 0.5) * invscale)\n",
    "weights = torch.where(j < xsize, weights, 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "id": "43e49804-2947-449f-8357-47e5919c9e1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([15, 4]),\n",
       " tensor([[0.5800, 0.0600, 0.0200, 0.1400],\n",
       "         [0.7400, 0.2200, 0.1800, 0.3000],\n",
       "         [0.9000, 0.3800, 0.3400, 0.4600],\n",
       "         [0.9400, 0.5400, 0.5000, 0.6200],\n",
       "         [0.7800, 0.7000, 0.6600, 0.7800],\n",
       "         [0.6200, 0.8600, 0.8200, 0.9400],\n",
       "         [0.4600, 0.9800, 0.9800, 0.9000],\n",
       "         [0.3000, 0.8200, 0.8600, 0.7400],\n",
       "         [0.1400, 0.6600, 0.7000, 0.5800],\n",
       "         [0.0000, 0.5000, 0.5400, 0.0000],\n",
       "         [0.0000, 0.3400, 0.3800, 0.0000],\n",
       "         [0.0000, 0.1800, 0.2200, 0.0000],\n",
       "         [0.0000, 0.0200, 0.0600, 0.0000],\n",
       "         [0.0000, 0.0000, 0.0000, 0.0000],\n",
       "         [0.0000, 0.0000, 0.0000, 0.0000]]))"
      ]
     },
     "execution_count": 203,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights.shape, weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "id": "05c9f3bf-42bb-4c58-998e-246b60909f50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([5.4600, 6.2600, 6.2600, 5.4600])"
      ]
     },
     "execution_count": 204,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_weights = weights.sum(dim=0)\n",
    "total_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "id": "75fe57c9-9931-4ae9-9989-b4a4e8951845",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.1062, 0.0096, 0.0032, 0.0256],\n",
       "        [0.1355, 0.0351, 0.0288, 0.0549],\n",
       "        [0.1648, 0.0607, 0.0543, 0.0842],\n",
       "        [0.1722, 0.0863, 0.0799, 0.1136],\n",
       "        [0.1429, 0.1118, 0.1054, 0.1429],\n",
       "        [0.1136, 0.1374, 0.1310, 0.1722],\n",
       "        [0.0842, 0.1565, 0.1565, 0.1648],\n",
       "        [0.0549, 0.1310, 0.1374, 0.1355],\n",
       "        [0.0256, 0.1054, 0.1118, 0.1062],\n",
       "        [0.0000, 0.0799, 0.0863, 0.0000],\n",
       "        [0.0000, 0.0543, 0.0607, 0.0000],\n",
       "        [0.0000, 0.0288, 0.0351, 0.0000],\n",
       "        [0.0000, 0.0032, 0.0096, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.0000, 0.0000, 0.0000, 0.0000]])"
      ]
     },
     "execution_count": 205,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights = weights / total_weights\n",
    "weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "id": "1582011a-ff7f-4103-85c0-ec6f9f6b4a86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([ 0,  3,  9, 16]), tensor([ 9, 13, 13,  9]))"
      ]
     },
     "execution_count": 195,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xmin, xsize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "id": "42facaeb-b8d5-4913-9fdb-449bae25eb18",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_idx = torch.arange(n, device=input_tensor.device).view(n, 1, 1, 1)\n",
    "c_idx = torch.arange(c, device=input_tensor.device).view(1, c, 1, 1)\n",
    "in_y = torch.arange(in_h, device=input_tensor.device).view((1, 1, in_h, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "id": "3cb1ca94-6197-4d64-854b-f00af116f7ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "xmin_idx = xmin.view(1, 1, 1, out_w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "id": "a8f1f006-ba97-4c93-9375-2ed43f39b1cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "in_tensors = [input_tensor[n_idx, c_idx, in_y, torch.clamp(xmin_idx + k, max=in_w - 1)] for k in range(max_interp_size.item())]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "id": "4b65cd02-7713-4d07-bfc9-6125637d138d",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_tensors = weights.unbind(dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "id": "440a04e0-d9a9-4d71-a3c6-4c274ede83a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial, reduce\n",
    "from typing import Callable, cast, Iterable, List, Optional, Tuple, Union\n",
    "from torch import sym_float, sym_int, Tensor\n",
    "\n",
    "\n",
    "def _sum_tensors(ts: Iterable[Tensor]) -> Tensor:\n",
    "    return reduce(torch.add, ts)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 214,
   "id": "92c85580-6501-4b42-b976-c66c507015ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = _sum_tensors(in_t * w_t for (in_t, w_t) in zip(in_tensors, w_tensors))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "id": "d9504e43-b3fd-42ce-846b-3ec9c70f35f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 3, 24, 4])"
      ]
     },
     "execution_count": 215,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af91ea2b-c736-4557-aaec-b25911abe724",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "id": "b20a44ae-6d63-46e8-97d4-f084113e7934",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import functional as F\n",
    "\n",
    "expected = F.interpolate(input_tensor, size=(out_h, out_w), mode=\"bilinear\", align_corners=align_corners, antialias=True)\n",
    "\n",
    "print(expected.shape)\n",
    "print(expected[0, 0, :5, :])\n",
    "print(output[0, 0, :5, :])\n",
    "\n",
    "torch.testing.assert_close(expected, output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f23676f1-508d-474b-8ad9-d973bba72b9a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}