ch02_5_JAX의_난수.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyPC7c004f6B5MnLBeqri5MY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/yblee110/jax-flax-book/blob/main/ch02_5_JAX%EC%9D%98_%EB%82%9C%EC%88%98.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install jax==0.4.26\n",
        "import jax\n",
        "print(jax.__version__)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T7j4RrjHiLjN",
        "outputId": "9927d039-080d-4a57-c282-a656c59eb299"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: jax==0.4.26 in /usr/local/lib/python3.10/dist-packages (0.4.26)\n",
            "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (0.2.0)\n",
            "Requirement already satisfied: numpy>=1.22 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (1.25.2)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (3.3.0)\n",
            "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (1.11.4)\n",
            "0.4.26\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3HtAubI6nmhh"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "np.random.seed(0)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def print_truncated_random_state():\n",
        "    \"\"\"출력에 문제생기지 않게 상태의 일부만 보여준다.\"\"\"\n",
        "    full_random_state = np.random.get_state()\n",
        "    print(str(full_random_state)[:460], '...')\n",
        "\n",
        "\n",
        "print_truncated_random_state()\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8yWmVlljnvhc",
        "outputId": "39949bbb-6392-47f0-bd0f-42429f103e38"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "('MT19937', array([         0,          1, 1812433255, 1900727105, 1208447044,\n",
            "       2481403966, 4042607538,  337614300, 3232553940, 1018809052,\n",
            "       3202401494, 1775180719, 3192392114,  594215549,  184016991,\n",
            "        829906058,  610491522, 3879932251, 3139825610,  297902587,\n",
            "       4075895579, 2943625357, 3530655617, 1423771745, 2135928312,\n",
            "       2891506774, 1066338622,  135451537,  933040465, 2759011858,\n",
            "       2273819758, 3545703099, 2516396728, 127 ...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "np.random.seed(0)\n",
        "print_truncated_random_state()\n",
        "_ = np.random.uniform()\n",
        "print_truncated_random_state()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vwUpfGoUoA5y",
        "outputId": "a04192fc-fd88-4c1c-f5a2-5a88f584c6e1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "('MT19937', array([         0,          1, 1812433255, 1900727105, 1208447044,\n",
            "       2481403966, 4042607538,  337614300, 3232553940, 1018809052,\n",
            "       3202401494, 1775180719, 3192392114,  594215549,  184016991,\n",
            "        829906058,  610491522, 3879932251, 3139825610,  297902587,\n",
            "       4075895579, 2943625357, 3530655617, 1423771745, 2135928312,\n",
            "       2891506774, 1066338622,  135451537,  933040465, 2759011858,\n",
            "       2273819758, 3545703099, 2516396728, 127 ...\n",
            "('MT19937', array([2443250962, 1093594115, 1878467924, 2709361018, 1101979660,\n",
            "       3904844661,  676747479, 2085143622, 1056793272, 3812477442,\n",
            "       2168787041,  275552121, 2696932952, 3432054210, 1657102335,\n",
            "       3518946594,  962584079, 1051271004, 3806145045, 1414436097,\n",
            "       2032348584, 1661738718, 1116708477, 2562755208, 3176189976,\n",
            "        696824676, 2399811678, 3992505346,  569184356, 2626558620,\n",
            "        136797809, 4273176064,  296167901, 343 ...\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "np.random.seed(0)\n",
        "print(np.random.uniform(size=3))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2hW-4SPhoFoz",
        "outputId": "89de8f3d-f2d3-40ab-ceea-e0b85992348f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0.5488135  0.71518937 0.60276338]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "np.random.seed(0)\n",
        "print(\"individually:\", np.stack([np.random.uniform() for _ in range(3)]))\n",
        "\n",
        "np.random.seed(0)\n",
        "print(\"all at once: \", np.random.uniform(size=3))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B78w3A94oIbt",
        "outputId": "d0fd20ab-7e4c-45e4-f1af-a701bcdef9ca"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "individually: [0.5488135  0.71518937 0.60276338]\n",
            "all at once:  [0.5488135  0.71518937 0.60276338]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "np.random.seed(0)\n",
        "\n",
        "def bar(): return np.random.uniform()\n",
        "def baz(): return np.random.uniform()\n",
        "\n",
        "\n",
        "def foo(): return bar() + 2 * baz()\n",
        "\n",
        "\n",
        "print(foo())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Hc_3t8CPoKWt",
        "outputId": "93b81223-f939-4bca-f23b-68b670055deb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.9791922366721637\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 새로운 버전\n",
        "from numpy.random import default_rng\n",
        "rng = default_rng(seed=0)\n",
        "vals = rng.standard_normal(10)\n",
        "more_vals = rng.standard_normal(10)\n",
        "\n",
        "# 이전버전\n",
        "from numpy import random\n",
        "vals = random.standard_normal(10)\n",
        "more_vals = random.standard_normal(10)\n"
      ],
      "metadata": {
        "id": "GNI_jjU_oNgE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from numpy.random import SeedSequence, default_rng\n",
        "ss = SeedSequence(12345)\n",
        "child_seeds = ss.spawn(10)\n",
        "streams = [default_rng(s) for s in child_seeds]\n",
        "print(streams)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0PBYgyO6oUSV",
        "outputId": "a0768adb-a09b-496f-e7b7-28c7bfdabb48"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[Generator(PCG64) at 0x7A03A6111000, Generator(PCG64) at 0x7A03A61111C0, Generator(PCG64) at 0x7A03A6110F20, Generator(PCG64) at 0x7A03A61118C0, Generator(PCG64) at 0x7A03A61110E0, Generator(PCG64) at 0x7A03A6111700, Generator(PCG64) at 0x7A03A6111620, Generator(PCG64) at 0x7A03A61117E0, Generator(PCG64) at 0x7A03A61119A0, Generator(PCG64) at 0x7A03A6111C40]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "from numpy.random import SeedSequence, default_rng\n",
        "ss = SeedSequence(12345)\n",
        "seeds = ss.spawn(2)\n",
        "stream = [default_rng(s) for s in seeds]\n",
        "\n",
        "def bar(): return stream[0].uniform()\n",
        "def baz(): return stream[1].uniform()\n",
        "\n",
        "def foo(): return bar() + 2 * baz()\n",
        "\n",
        "print(foo())\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3Ilv61CeoXGt",
        "outputId": "d66adbff-8057-433d-cf1c-f8c702628f76"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.6241496684412051\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from jax import random\n",
        "key = random.PRNGKey(42)\n",
        "print(key)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EhDqfygvodUA",
        "outputId": "202e1019-a037-4fbc-bdfb-0bdab666251a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[ 0 42]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(random.normal(key))\n",
        "print(random.normal(key))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mCF7_reLofpD",
        "outputId": "25d196ab-25b4-48ad-9ffa-81b340a72658"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "-0.18471177\n",
            "-0.18471177\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"old key\", key)\n",
        "new_key, subkey = random.split(key)\n",
        "del key  # 오래된 키는 지워버리며 나중에라도 사용ㅇ하지 않습니다..\n",
        "normal_sample = random.normal(subkey)\n",
        "print(r\"    \\---SPLIT --> new key   \", new_key)\n",
        "print(r\"             \\--> new subkey\", subkey, \"--> normal\", normal_sample)\n",
        "del subkey  # 서브키도 사용후에 제거해야 합니다.\n",
        "key = new_key  # 만약에 한번 이 키를 다시 생성해야 한다면 new_key는 키로 사용됩니다."
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O9eTfgXToizQ",
        "outputId": "e28dd160-a1d7-4346-a9d9-37242bba3bc6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "old key [ 0 42]\n",
            "    \\---SPLIT --> new key    [2465931498 3679230171]\n",
            "             \\--> new subkey [255383827 267815257] --> normal 1.3694694\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "key, subkey = random.split(key)"
      ],
      "metadata": {
        "id": "cNnWNQa9ok2A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "key = random.PRNGKey(42)\n",
        "subkeys = random.split(key, 3)\n",
        "sequence = np.stack([random.normal(subkey) for subkey in subkeys])\n",
        "print(\"individually:\", sequence)\n",
        "key = random.PRNGKey(42)\n",
        "print(\"all at once: \", random.normal(key, shape=(3,)))\n"
      ],
      "metadata": {
        "id": "Px_18uMGo-iN",
        "outputId": "f616939e-cfba-4ae1-c4b6-5dc6ca89048d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "individually: [-0.04838832  0.10796154 -1.2226542 ]\n",
            "all at once:  [ 0.18693547 -1.2806505  -1.5593132 ]\n"
          ]
        }
      ]
    }
  ]
}