docs/source/notebooks/nan_grad.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "237dd11d-c129-4a43-8731-461bf1c36663",
   "metadata": {},
   "source": [
    "# Distinguishing between 0 and NaN gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64da0598-b7a7-4856-a199-b29e63db38d2",
   "metadata": {},
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pytorch/maskedtensor/blob/main/docs/source/notebooks/nan_grad.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "275907cb-25b2-4b24-bd5c-f48f5dfc7b7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from maskedtensor import masked_tensor\n",
    "from maskedtensor import as_masked_tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf1f8eb-6eef-415a-9351-b36dcbe0a1f8",
   "metadata": {},
   "source": [
    "## Resolving Issues"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6f10ac2-5232-4cf1-8b7e-e772a973f3d0",
   "metadata": {},
   "source": [
    "One issue that vanilla tensors run into is the inability to distinguish between gradients that are not defined (nan) vs. gradients that are actually 0.\n",
    "\n",
    "Below, by way of example, we show several different issues where `torch.Tensor` falls short and `MaskedTensor` can resolve and/or work around the NaN gradient problem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fac63235-de62-4b40-9f57-e2726d7897c8",
   "metadata": {
    "tags": []
   },
   "source": [
    "### [PyTorch Issue 10729 - torch.where](https://github.com/pytorch/pytorch/issues/10729)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cec9c06-8f3f-4515-925d-ed8282ceaf42",
   "metadata": {},
   "source": [
    "**PyTorch result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "329bbdb2-e3e3-448b-976f-5fdecb950849",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y: tensor([4.5400e-05, 6.7379e-03, 1.0000e+00, 1.0000e+00, 1.0000e+00, 1.0000e+00,\n",
      "        1.0000e+00, 1.0000e+00, 1.0000e+00, 1.0000e+00, 1.0000e+00],\n",
      "       grad_fn=<SWhereBackward0>)\n",
      "x.grad: tensor([4.5400e-05, 6.7379e-03, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
      "        0.0000e+00, 0.0000e+00, 0.0000e+00,        nan,        nan])\n",
      "y.grad: None\n"
     ]
    }
   ],
   "source": [
    "# This behavior underlies the fix to clamp, which uses where in its derivative\n",
    "x = torch.tensor([-10., -5, 0, 5, 10, 50, 60, 70, 80, 90, 100], requires_grad=True)\n",
    "y = torch.where(x < 0, torch.exp(x), torch.ones_like(x))\n",
    "print(\"y:\", y)\n",
    "y.sum().backward()\n",
    "print(\"x.grad:\", x.grad)\n",
    "print(\"y.grad:\", y.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99fc0741-ae9c-4489-8273-f2f7a0c6557f",
   "metadata": {},
   "source": [
    "**MaskedTensor result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "396f6b1e-4e84-40ce-ba8d-989dad6cb2d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mx.grad:  masked_tensor(\n",
      "  [  0.0000,   0.0067,       --,       --,       --,       --,       --,       --,       --,       --,       --]\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "x = torch.tensor([-10., -5, 0, 5, 10, 50, 60, 70, 80, 90, 100], requires_grad=True)\n",
    "mask = x < 0\n",
    "mx = masked_tensor(x, mask, requires_grad=True)\n",
    "my = masked_tensor(torch.ones_like(x), ~mask, requires_grad=True)\n",
    "y = torch.where(mask, torch.exp(mx), my)\n",
    "s = y.sum()\n",
    "s.backward()\n",
    "# Gradient is only provided to selected subset.\n",
    "# Effectively this changes the gradient of where to mask out elements instead\n",
    "# of setting them to zero.\n",
    "print(\"mx.grad: \", mx.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2cedd39-bb73-4f5a-932d-4aa776f0edbb",
   "metadata": {},
   "source": [
    "The gradient here is only provided to the selected subset. Effectively, this changes the gradient of where to mask out elements instead of setting them to zero."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "235941eb-5544-480d-90dc-1f5caf754b97",
   "metadata": {},
   "source": [
    "### [PyTorch Issue 52248 - another torch.where](https://github.com/pytorch/pytorch/issues/52248)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10a825d0-725f-4e20-abbb-34bcdb32feea",
   "metadata": {},
   "source": [
    "**PyTorch result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8831cb6c-5e89-4b0f-817f-3d749408144a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1., grad_fn=<SWhereBackward0>)\n",
      "(tensor(nan),)\n"
     ]
    }
   ],
   "source": [
    "# A more recent incarnation specific to where of this\n",
    "# https://github.com/pytorch/pytorch/issues/52248\n",
    "\n",
    "a = torch.randn((), requires_grad=True)\n",
    "b = torch.tensor(False)\n",
    "c = torch.ones(())\n",
    "\n",
    "print(torch.where(b, a/0, c))\n",
    "print(torch.autograd.grad(torch.where(b, a/0, c), a))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02af6a11-daf6-4be5-9120-76fe694524a7",
   "metadata": {},
   "source": [
    "**MaskedTensor result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "df8f0244-2f73-4ef1-9c6c-d176edda0557",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "masked_tensor(  1.0000, True)\n",
      "(masked_tensor(--, False),)\n"
     ]
    }
   ],
   "source": [
    "a = masked_tensor(torch.randn(()), torch.tensor(True), requires_grad=True)\n",
    "b = torch.tensor(False)\n",
    "c = torch.ones(())\n",
    "\n",
    "print(torch.where(b, a/0, c))\n",
    "print(torch.autograd.grad(torch.where(b, a/0, c), a))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3a03ca1-37c9-40fe-859e-bb3d7001c922",
   "metadata": {},
   "source": [
    "### [PyTorch Issue 67180 - torch.nansum and torch.nanmean](https://github.com/pytorch/pytorch/issues/67180)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e9c88b6-1714-41a7-80c3-ea8f9378902a",
   "metadata": {},
   "source": [
    "**PyTorch result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "66b53eb6-1220-4fef-9d35-4379a1c4ebbd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(nan)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.tensor([1., 2., float('nan')])\n",
    "b = torch.tensor(1.0, requires_grad=True)\n",
    "c = a * b\n",
    "c1 = torch.nansum(c)  # or torch.nanmean\n",
    "\n",
    "bgrad1, = torch.autograd.grad(c1, b, retain_graph=True)\n",
    "bgrad1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d01260-ae4d-47e1-84de-d5d0ac9d5ebb",
   "metadata": {},
   "source": [
    "**MaskedTensor result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c0ec6400-cfbf-4c36-a342-ccee22d93013",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "masked_tensor(  3.0000, True)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.tensor([1., 2., float('nan')])\n",
    "b = torch.tensor(1.0, requires_grad=True)\n",
    "ma = masked_tensor(a, ~torch.isnan(a))\n",
    "c = ma * b\n",
    "c1 = torch.sum(c)  # or torch.nanmean\n",
    "\n",
    "bgrad1, = torch.autograd.grad(c1, b, retain_graph=True)\n",
    "bgrad1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70ed85c7-a0e8-4981-a6c9-43153de1dc86",
   "metadata": {},
   "source": [
    "### [PyTorch Issue 4132 - when using mask, x/0 yields NaN grad](https://github.com/pytorch/pytorch/issues/4132)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "652c8a5a-3fd0-4b48-add4-b258249ab739",
   "metadata": {},
   "source": [
    "**PyTorch result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "75427dc5-f5e2-46e5-90c0-62fca729a266",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([nan, 1.])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.tensor([1., 1.], requires_grad=True)\n",
    "div = torch.tensor([0., 1.])\n",
    "y = x/div # => y is [inf, 1]\n",
    "\n",
    "mask = (div != 0) # => mask is [0, 1]\n",
    "loss = y[mask]\n",
    "loss.backward()\n",
    "\n",
    "x.grad # grad is [nan, 1], but expected [0, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "141c06c8-7555-45ad-9fd1-7fb64345ea1e",
   "metadata": {},
   "source": [
    "**MaskedTensor result**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2a02d643-0531-4d18-bca6-8b4e535c7223",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "masked_tensor(\n",
       "  [      --,   1.0000]\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.tensor([1., 1.], requires_grad=True)\n",
    "div = torch.tensor([0., 1.])\n",
    "y = x/div # => y is [inf, 1]\n",
    "\n",
    "mask = (div != 0) # => mask is [0, 1]\n",
    "loss = as_masked_tensor(y, mask)\n",
    "# We could add autograd support for indexing here instead of using sum\n",
    "loss = loss.sum()\n",
    "loss.backward()\n",
    "\n",
    "x.grad # grad is [nan, 1], but expected [0, 1]"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('pytorch_env')",
   "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.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "d0a07e0fa82e52b776976e55e335decf2d0ac48acfa03bb0e56e7f3ca52a96d6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}