torch.multinomial picking 0 prob element.ipynb

torch-multinomial-picking-0-prob-element.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNuUrORzRAf2LTkHYo4jw+W",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/myazdani/75b5152d257c25ec331e04e6247b2e39/torch-multinomial-picking-0-prob-element.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Originally reported: https://github.com/pytorch/pytorch/issues/13867"
      ],
      "metadata": {
        "id": "ALTYChFra4Qp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "KQqo-CjGadP4"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1)\n",
        "# top k sampling from normalized logits\n",
        "k=3\n",
        "for i in range(100):\n",
        "    logits = torch.randn(32,50257)\n",
        "    probs = F.softmax(logits, dim=-1)\n",
        "    samp_probs = probs.clone()\n",
        "    indices_to_remove = samp_probs < torch.topk(samp_probs, k)[0][..., -1, None]\n",
        "    samp_probs[indices_to_remove] = 0\n",
        "    samp_probs.div_(samp_probs.sum(1).unsqueeze(1))\n",
        "    next_tokens = samp_probs.multinomial(1)\n",
        "    next_logprobs = samp_probs.gather(1, next_tokens.view(-1, 1)).log()\n",
        "    if next_logprobs.isinf().any().item():\n",
        "        print(f\"trial {i} failed\")\n",
        "        break"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UlA4GlMjafTa",
        "outputId": "e762f465-38f0-4614-dabc-1922cab185c7"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "trial 2 failed\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "-UgdKiFzarCJ"
      },
      "execution_count": 2,
      "outputs": []
    }
  ]
}