JAX_Smooth_Smith_Waterman.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "JAX Smooth Smith-Waterman",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMc5lSroqkfn+w0wXcnpabU",
      "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/github/shu65/blog-jax-notebook/blob/main/JAX_Smooth_Smith_Waterman.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kt_1erhlCW6h"
      },
      "source": [
        "import numpy as np\n",
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EqxH7QZeCnve"
      },
      "source": [
        "np.random.seed(0)\n",
        "seq_1_len = 100\n",
        "seq_2_len = 150\n",
        "score_matrix_np = np.random.random((seq_1_len, seq_2_len))\n",
        "score_matrix_jnp = jnp.array(score_matrix_np)"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e5vj8F45Cq2Q",
        "outputId": "da800e42-9ab0-4b4f-cfe1-377a78866fb2"
      },
      "source": [
        "def sw_np(NINF=-1e30):\n",
        "    \n",
        "    def _logsumexp(y, axis):\n",
        "        y = np.maximum(y,NINF)\n",
        "        return y.max(axis) + np.log(np.sum(np.exp(y - y.max(axis, keepdims=True)), axis=axis))\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        real_a, real_b = lengths\n",
        "        hij = np.full((real_a + 1, real_b + 1), fill_value=NINF, dtype=np.float32)\n",
        "        for i in range(real_a):\n",
        "            for j in range(real_b):\n",
        "                s = score_matrix[i, j]\n",
        "                m = hij[i, j] + s\n",
        "                g0 = hij[i + 1, j] + gap\n",
        "                g1 = hij[i, j + 1] + gap\n",
        "\n",
        "                h = np.stack([m, g0, g1, s], -1)\n",
        "                hij[i + 1, j + 1] = _soft_maximum(h, temp=temp, axis=-1)\n",
        "        hij = hij[1:, 1:]\n",
        "        score = _soft_maximum(hij, temp=temp)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "my_sw_func = sw_np()\n",
        "%time score = my_sw_func(score_matrix_np, (seq_1_len, seq_2_len))\n",
        "print(score)"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 735 ms, sys: 4 ms, total: 739 ms\n",
            "Wall time: 743 ms\n",
            "232.31179809570312\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Min2FfsqCvBg"
      },
      "source": [
        "def sw_v0(NINF=-1e30):\n",
        "    \n",
        "    def _logsumexp(y, axis):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return jax.nn.logsumexp(y, axis=axis)\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        real_a, real_b = lengths\n",
        "        hij = jnp.full((real_a + 1, real_b + 1), fill_value=NINF, dtype=jnp.float32)\n",
        "        for i in range(real_a):\n",
        "            for j in range(real_b):\n",
        "                s = score_matrix[i, j]\n",
        "                m = hij[i, j] + s\n",
        "                g0 = hij[i + 1, j] + gap\n",
        "                g1 = hij[i, j + 1] + gap\n",
        "                h = jnp.stack([m, g0, g1, s], -1)\n",
        "                hij = hij.at[i + 1, j + 1].set(_soft_maximum(h, -1))\n",
        "        hij = hij[1:, 1:]\n",
        "        score = _soft_maximum(hij)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "# this is too slow\n",
        "#my_sw_func = sw_v0()\n",
        "#print(\"jax default first call\")\n",
        "#%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "#print(\"jax default second call\")\n",
        "#%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "#print(score)\n",
        "#print()\n",
        "\n",
        "#my_sw_func = jax.jit(sw_v0())\n",
        "#print(\"jax jit first call\")\n",
        "#%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "#print(\"jax jit second call\")\n",
        "#%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "#print(score)\n",
        "#print()"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UyoFB4wGCyzP",
        "outputId": "6c4d5013-0e9e-4952-e3b8-4740ffa4e4eb"
      },
      "source": [
        "def sw_v1(unroll=2, NINF=-1e30):\n",
        "    \n",
        "    def _rotate(score_matrix):\n",
        "        a,b = score_matrix.shape\n",
        "        n,m = (a+b-1),(a+b)//2\n",
        "        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]\n",
        "        i,j = (br-ar)+(a-1),(ar+br)//2\n",
        "        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)\n",
        "        reverse_idx = (i, j)\n",
        "        return rotated_score_matrix, reverse_idx\n",
        "    \n",
        "    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):\n",
        "        return rotated_dp_matrix[reverse_idx]\n",
        "\n",
        "    def _logsumexp(y, axis):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return jax.nn.logsumexp(y, axis=axis)\n",
        "\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _step(prev, gap_cell_condition, rotated_score_matrix, gap, temp):\n",
        "        h2,h1 = prev   # previous two rows of scoring (hij) mtx\n",
        "        h1_T = jax.lax.cond(\n",
        "            gap_cell_condition,\n",
        "            lambda x: jnp.pad(x[:-1], [1,0], constant_values=(NINF,NINF)),\n",
        "            lambda x: jnp.pad(x[1:], [0,1], constant_values=(NINF,NINF)),\n",
        "            h1,\n",
        "        )\n",
        "\n",
        "        a = h2 + rotated_score_matrix\n",
        "        g0 = h1 + gap\n",
        "        g1 = h1_T + gap\n",
        "        s = rotated_score_matrix\n",
        "\n",
        "        h0 = jnp.stack([a, g0, g1, s], -1)\n",
        "        h0 = _soft_maximum(h0, temp, -1)\n",
        "        return (h1,h0), h0\n",
        "\n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        rotated_score_matrix, reverse_idx = _rotate(score_matrix)\n",
        "        \n",
        "        a,b = score_matrix.shape\n",
        "        n,m = rotated_score_matrix.shape\n",
        "        \n",
        "        gap_cell_condition = (jnp.arange(n)+a%2)%2\n",
        "        prev = (jnp.full(m, NINF), jnp.full(m, NINF))\n",
        "        rotated_hij = []\n",
        "        for i in range(n):\n",
        "            prev, h = _step(prev, gap_cell_condition[i], rotated_score_matrix[i], gap, temp)\n",
        "            rotated_hij.append(h)\n",
        "        rotated_hij = jnp.stack(rotated_hij)\n",
        "        hij = _rotate_in_reverse(rotated_hij, reverse_idx)\n",
        "        score = _soft_maximum(hij, temp=temp)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "my_sw_func = sw_v1()\n",
        "print(\"jax default first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax default second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()\n",
        "\n",
        "my_sw_func = jax.jit(sw_v1())\n",
        "print(\"jax jit first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax jit second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "jax default first call\n",
            "CPU times: user 17.7 s, sys: 177 ms, total: 17.8 s\n",
            "Wall time: 17.8 s\n",
            "jax default second call\n",
            "CPU times: user 17.6 s, sys: 153 ms, total: 17.7 s\n",
            "Wall time: 17.7 s\n",
            "232.31181\n",
            "\n",
            "jax jit first call\n",
            "CPU times: user 2min 20s, sys: 715 ms, total: 2min 21s\n",
            "Wall time: 2min 20s\n",
            "jax jit second call\n",
            "CPU times: user 1.98 ms, sys: 0 ns, total: 1.98 ms\n",
            "Wall time: 1.81 ms\n",
            "232.31181\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NRyOeUCLDkGo",
        "outputId": "5a2f695e-4bb3-4e49-8d98-31f179081d03"
      },
      "source": [
        "def sw_v2(unroll=2, NINF=-1e30):\n",
        "\n",
        "    def _rotate(score_matrix):\n",
        "        a,b = score_matrix.shape\n",
        "        n,m = (a+b-1),(a+b)//2\n",
        "        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]\n",
        "        i,j = (br-ar)+(a-1),(ar+br)//2\n",
        "        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)\n",
        "        reverse_idx = (i, j)\n",
        "        return rotated_score_matrix, reverse_idx\n",
        "\n",
        "    def _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp):\n",
        "        def scan_f(prev, scan_xs):\n",
        "            h2, h1 = prev\n",
        "            h1_T = jax.lax.cond(\n",
        "                scan_xs[\"gap_cell_condition\"],\n",
        "                lambda x: jnp.pad(x[:-1], [1,0], constant_values=(NINF,NINF)),\n",
        "                lambda x: jnp.pad(x[1:], [0,1], constant_values=(NINF,NINF)),\n",
        "                h1,\n",
        "            )\n",
        "            a = h2 + scan_xs[\"rotated_score_matrix\"]\n",
        "            g0 = h1 + gap\n",
        "            g1 = h1_T + gap\n",
        "            s = scan_xs[\"rotated_score_matrix\"]\n",
        "\n",
        "            h0 = jnp.stack([a, g0, g1, s], -1)\n",
        "            h0 = _soft_maximum(h0, temp, -1)\n",
        "            return (h1,h0), h0\n",
        "        \n",
        "        a,b = score_matrix.shape\n",
        "        n,m = rotated_score_matrix.shape\n",
        "\n",
        "        scan_xs = {\n",
        "            \"rotated_score_matrix\": rotated_score_matrix,\n",
        "            \"gap_cell_condition\": (jnp.arange(n)+a%2)%2\n",
        "        }\n",
        "        scan_init = (jnp.full(m, NINF), jnp.full(m, NINF))\n",
        "        return scan_f, scan_xs, scan_init\n",
        "\n",
        "    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):\n",
        "        return rotated_dp_matrix[reverse_idx]\n",
        "\n",
        "    def _logsumexp(y, axis):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return jax.nn.logsumexp(y, axis=axis)\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _get_prev_gap_cell_score(cond, true, false): \n",
        "        return cond*true + (1-cond)*false\n",
        "    \n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        rotated_score_matrix, reverse_idx = _rotate(score_matrix)\n",
        "        scan_f, scan_xs, scan_init = _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp)\n",
        "        rotated_hij = jax.lax.scan(scan_f, scan_init, scan_xs, unroll=unroll)[-1]\n",
        "        hij = _rotate_in_reverse(rotated_hij, reverse_idx)\n",
        "        score = _soft_maximum(hij, temp, axis=None)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "my_sw_func = sw_v2()\n",
        "print(\"jax default first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax default second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()\n",
        "\n",
        "my_sw_func = jax.jit(sw_v2())\n",
        "print(\"jax jit first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax jit second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "jax default first call\n",
            "CPU times: user 739 ms, sys: 18 ms, total: 757 ms\n",
            "Wall time: 758 ms\n",
            "jax default second call\n",
            "CPU times: user 666 ms, sys: 1.98 ms, total: 668 ms\n",
            "Wall time: 671 ms\n",
            "232.31181\n",
            "\n",
            "jax jit first call\n",
            "CPU times: user 1 s, sys: 5.01 ms, total: 1.01 s\n",
            "Wall time: 1.01 s\n",
            "jax jit second call\n",
            "CPU times: user 339 µs, sys: 989 µs, total: 1.33 ms\n",
            "Wall time: 1.14 ms\n",
            "232.31181\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "26PDRnNZWNiC",
        "outputId": "a8947e29-80ac-4bf3-cbda-e1fd2bbf2262"
      },
      "source": [
        "def sw_v3(unroll=2, NINF=-1e30):\n",
        "\n",
        "    def _rotate(score_matrix):\n",
        "        a,b = score_matrix.shape\n",
        "        n,m = (a+b-1),(a+b)//2\n",
        "        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]\n",
        "        i,j = (br-ar)+(a-1),(ar+br)//2\n",
        "        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)\n",
        "        reverse_idx = (i, j)\n",
        "        return rotated_score_matrix, reverse_idx\n",
        "\n",
        "    def _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp):\n",
        "        def scan_f(prev, scan_xs):\n",
        "            h2, h1 = prev\n",
        "            h1_T = _get_prev_gap_cell_score(\n",
        "                scan_xs[\"gap_cell_condition\"],\n",
        "                jnp.pad(h1[:-1], [1,0], constant_values=(NINF,NINF)),\n",
        "                jnp.pad(h1[1:], [0,1], constant_values=(NINF,NINF)),\n",
        "            )\n",
        "            a = h2 + scan_xs[\"rotated_score_matrix\"]\n",
        "            g0 = h1 + gap\n",
        "            g1 = h1_T + gap\n",
        "            s = scan_xs[\"rotated_score_matrix\"]\n",
        "\n",
        "            h0 = jnp.stack([a, g0, g1, s], -1)\n",
        "            h0 = _soft_maximum(h0, temp, -1)\n",
        "            return (h1,h0), h0\n",
        "        \n",
        "        a,b = score_matrix.shape\n",
        "        n,m = rotated_score_matrix.shape\n",
        "\n",
        "        scan_xs = {\n",
        "            \"rotated_score_matrix\": rotated_score_matrix,\n",
        "            \"gap_cell_condition\": (jnp.arange(n)+a%2)%2\n",
        "        }\n",
        "        scan_init = (jnp.full(m, NINF), jnp.full(m, NINF))\n",
        "        return scan_f, scan_xs, scan_init\n",
        "\n",
        "    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):\n",
        "        return rotated_dp_matrix[reverse_idx]\n",
        "\n",
        "    def _logsumexp(y, axis):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return jax.nn.logsumexp(y, axis=axis)\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _get_prev_gap_cell_score(cond, true, false): \n",
        "        return cond*true + (1-cond)*false\n",
        "    \n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        rotated_score_matrix, reverse_idx = _rotate(score_matrix)\n",
        "        scan_f, scan_xs, scan_init = _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp)\n",
        "        rotated_hij = jax.lax.scan(scan_f, scan_init, scan_xs, unroll=unroll)[-1]\n",
        "        hij = _rotate_in_reverse(rotated_hij, reverse_idx)\n",
        "        score = _soft_maximum(hij, temp, axis=None)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "my_sw_func = sw_v3()\n",
        "print(\"jax default first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax default second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()\n",
        "\n",
        "my_sw_func = jax.jit(sw_v3())\n",
        "print(\"jax jit first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax jit second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jax default first call\n",
            "CPU times: user 599 ms, sys: 1.99 ms, total: 601 ms\n",
            "Wall time: 608 ms\n",
            "jax default second call\n",
            "CPU times: user 599 ms, sys: 3.02 ms, total: 602 ms\n",
            "Wall time: 607 ms\n",
            "232.31181\n",
            "\n",
            "jax jit first call\n",
            "CPU times: user 940 ms, sys: 2.01 ms, total: 942 ms\n",
            "Wall time: 947 ms\n",
            "jax jit second call\n",
            "CPU times: user 4.9 ms, sys: 0 ns, total: 4.9 ms\n",
            "Wall time: 3.41 ms\n",
            "232.31181\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6Mvys2d1DmHw",
        "outputId": "142e23d7-ee40-4485-fc8e-29f472db70f8"
      },
      "source": [
        "def sw_v4(unroll=2, NINF=-1e30):\n",
        "    \n",
        "    def _make_mask(score_matrix, lengths):\n",
        "        a,b = score_matrix.shape\n",
        "        real_a, real_b = lengths\n",
        "        mask = (jnp.arange(a) < real_a)[:,None] * (jnp.arange(b) < real_b)[None,:]\n",
        "        return mask\n",
        "\n",
        "    def _rotate(score_matrix):\n",
        "        a,b = score_matrix.shape\n",
        "        n,m = (a+b-1),(a+b)//2\n",
        "        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]\n",
        "        i,j = (br-ar)+(a-1),(ar+br)//2\n",
        "        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)\n",
        "        reverse_idx = (i, j)\n",
        "        return rotated_score_matrix, reverse_idx\n",
        "\n",
        "    def _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp):\n",
        "        def scan_f(prev, scan_xs):\n",
        "            h2, h1 = prev\n",
        "            h1_T = _get_prev_gap_cell_score(\n",
        "                scan_xs[\"gap_cell_condition\"],\n",
        "                jnp.pad(h1[:-1], [1,0], constant_values=(NINF,NINF)),\n",
        "                jnp.pad(h1[1:], [0,1], constant_values=(NINF,NINF)),\n",
        "            )\n",
        "            a = h2 + scan_xs[\"rotated_score_matrix\"]\n",
        "            g0 = h1 + gap\n",
        "            g1 = h1_T + gap\n",
        "            s = scan_xs[\"rotated_score_matrix\"]\n",
        "\n",
        "            h0 = jnp.stack([a, g0, g1, s], -1)\n",
        "            h0 = _soft_maximum(h0, temp, -1)\n",
        "            return (h1,h0), h0\n",
        "        \n",
        "        a,b = score_matrix.shape\n",
        "        n,m = rotated_score_matrix.shape\n",
        "\n",
        "        scan_xs = {\n",
        "            \"rotated_score_matrix\": rotated_score_matrix,\n",
        "            \"gap_cell_condition\": (jnp.arange(n)+a%2)%2\n",
        "        }\n",
        "        scan_init = (jnp.full(m, NINF), jnp.full(m, NINF))\n",
        "        return scan_f, scan_xs, scan_init\n",
        "\n",
        "    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):\n",
        "        return rotated_dp_matrix[reverse_idx]\n",
        "\n",
        "    def _logsumexp(y, axis):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return jax.nn.logsumexp(y, axis=axis)\n",
        "\n",
        "    def _logsumexp_with_mask(y, axis, mask):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return y.max(axis) + jnp.log(jnp.sum(mask * jnp.exp(y - y.max(axis, keepdims=True)), axis=axis))\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _soft_maximum_with_mask(x, temp, mask, axis=None):\n",
        "        return temp*_logsumexp_with_mask(x/temp, axis, mask)\n",
        "\n",
        "    def _get_prev_gap_cell_score(cond, true, false): \n",
        "        return cond*true + (1-cond)*false\n",
        "    \n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        mask = _make_mask(score_matrix, lengths)\n",
        "        masked_score_matrix = score_matrix + NINF * (1 - mask)\n",
        "        rotated_score_matrix, reverse_idx = _rotate(masked_score_matrix)\n",
        "        scan_f, scan_xs, scan_init = _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp)\n",
        "        rotated_hij = jax.lax.scan(scan_f, scan_init, scan_xs, unroll=unroll)[-1]\n",
        "        hij = _rotate_in_reverse(rotated_hij, reverse_idx)\n",
        "        score = _soft_maximum_with_mask(hij, temp, mask=mask, axis=None)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "my_sw_func = sw_v4()\n",
        "print(\"jax default first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax default second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()\n",
        "\n",
        "my_sw_func = jax.jit(sw_v4())\n",
        "print(\"jax jit first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(\"jax jit second call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "print(score)\n",
        "print()"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jax default first call\n",
            "CPU times: user 577 ms, sys: 2 ms, total: 579 ms\n",
            "Wall time: 582 ms\n",
            "jax default second call\n",
            "CPU times: user 591 ms, sys: 0 ns, total: 591 ms\n",
            "Wall time: 595 ms\n",
            "232.31181\n",
            "\n",
            "jax jit first call\n",
            "CPU times: user 909 ms, sys: 1.99 ms, total: 911 ms\n",
            "Wall time: 912 ms\n",
            "jax jit second call\n",
            "CPU times: user 5.02 ms, sys: 0 ns, total: 5.02 ms\n",
            "Wall time: 3.53 ms\n",
            "232.31181\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CKoZHabCDrJq"
      },
      "source": [
        "seq_1_max_len = 100\n",
        "seq_2_max_len = 120\n",
        "num_pairs = 100\n",
        "\n",
        "batch_score_matrix_np = np.random.random((num_pairs, seq_1_len, seq_2_len))\n",
        "batch_lens_np = np.array([[np.random.choice([80,90,100]),np.random.choice([95,105,120])] for _ in range(num_pairs)])\n",
        "\n",
        "batch_score_matrix_jnp = jnp.array(batch_score_matrix_np)\n",
        "batch_lens_jnp = jnp.array(batch_lens_np)"
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dkflglkIDzdx",
        "outputId": "a617f05e-ac6a-4ffa-cdab-c7274ce5226f"
      },
      "source": [
        "def batch_sw_np(NINF=-1e30):\n",
        "    def _batch_sw(batch_score_matrix, batch_lengths, gap=0, temp=1.0):\n",
        "        n_batches = batch_score_matrix.shape[0]\n",
        "        sw_func = sw_np(NINF=NINF)\n",
        "        ret = [sw_func(batch_score_matrix[i], batch_lengths[i], gap=gap, temp=temp) \n",
        "               for i in range(n_batches)]\n",
        "        return np.array(ret)\n",
        "    return _batch_sw\n",
        "\n",
        "my_sw_func = batch_sw_np()\n",
        "print(\"batch np\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0)\n",
        "print(score)\n",
        "print()"
      ],
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "batch np\n",
            "CPU times: user 42.4 s, sys: 208 ms, total: 42.6 s\n",
            "Wall time: 42.4 s\n",
            "[117.34008026 105.93036652 117.13354492  98.28081512 104.79451752\n",
            " 117.04177856 111.91090393 107.09848785  98.45252991  97.85121155\n",
            "  96.686409    98.21727753  93.44289398 117.16491699  94.247612\n",
            "  96.47311401 106.85935211 108.07891083 111.64317322 104.01583862\n",
            " 108.43913269  97.67262268 106.19509888 109.37216187  98.36544037\n",
            "  93.17828369 104.95552063 107.21205139  93.5545578  104.81386566\n",
            "  97.9903183  117.33205414 117.09793091  96.43874359 104.06102753\n",
            " 104.08187866  94.45321655 118.06995392  98.21305084  93.45307159\n",
            "  98.7559433  108.422966    96.12508392 106.63827515 117.93927765\n",
            "  98.15522003 117.89131927  94.71160126 117.4874115   96.24068451\n",
            "  98.15233612 102.07528687 109.00423431 107.47882843  96.7345047\n",
            " 116.69099426 107.00982666 111.89899445 109.32804871 101.34351349\n",
            " 117.27951813 102.11739349  97.68030548  93.76876831 101.16060638\n",
            " 101.1467514   99.02960205 108.44194031  97.4659729  101.9356308\n",
            " 104.6388092   98.22956085  98.02523041  94.44136047 100.73300934\n",
            " 107.8992157   95.57440948 106.34977722  96.10211182 112.58492279\n",
            " 101.88635254  98.96472931 101.89956665  95.47772217 108.00675201\n",
            " 116.76008606  97.79943848  97.47626495 117.84597778 112.64025879\n",
            "  92.89571381 100.36632538  96.89391327 108.94417572 106.70108795\n",
            "  99.47631073  93.87617493 112.3595047   96.83467102 107.74719238]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7FTFtkFfD0wx",
        "outputId": "9272ebbe-b602-43d2-a414-e24f66c69d4b"
      },
      "source": [
        "def batch_sw_v0(NINF=-1e30):\n",
        "    def _batch_sw(batch_score_matrix, batch_lengths, gap=0, temp=1.0):\n",
        "        n_batches = batch_score_matrix.shape[0]\n",
        "        sw_func = jax.jit(sw_v4())\n",
        "        ret = [sw_func(batch_score_matrix[i], batch_lengths[i], gap, temp) \n",
        "               for i in range(n_batches)]\n",
        "        return jnp.array(ret)\n",
        "    return _batch_sw\n",
        "\n",
        "my_sw_func = batch_sw_v0()\n",
        "print(\"batch jax default first call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0)\n",
        "print(\"batch jax default second call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0)\n",
        "print(score)\n",
        "print()\n",
        "\n",
        "my_sw_func = jax.jit(batch_sw_v0())\n",
        "print(\"batch jax default first call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "print(\"batch jax default second call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "print(score)\n",
        "print()"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "batch jax default first call\n",
            "CPU times: user 1.31 s, sys: 13 ms, total: 1.33 s\n",
            "Wall time: 1.32 s\n",
            "batch jax default second call\n",
            "CPU times: user 1.3 s, sys: 5.02 ms, total: 1.3 s\n",
            "Wall time: 1.3 s\n",
            "[117.34008  105.93037  117.13354   98.280815 104.79451  117.04178\n",
            " 111.9109   107.09849   98.45252   97.8512    96.68642   98.21728\n",
            "  93.4429   117.164894  94.2476    96.47309  106.859344 108.07891\n",
            " 111.64319  104.01584  108.439125  97.67263  106.1951   109.37219\n",
            "  98.36543   93.17829  104.95553  107.21204   93.55455  104.81387\n",
            "  97.9903   117.332054 117.09793   96.438736 104.06102  104.08189\n",
            "  94.45322  118.06995   98.21306   93.453064  98.755936 108.42296\n",
            "  96.125084 106.638275 117.939285  98.15523  117.89131   94.7116\n",
            " 117.48742   96.24068   98.15234  102.07528  109.00425  107.478836\n",
            "  96.7345   116.69099  107.00981  111.89899  109.32806  101.34352\n",
            " 117.27952  102.117386  97.68031   93.76876  101.16061  101.14676\n",
            "  99.0296   108.44193   97.46598  101.93562  104.638794  98.22956\n",
            "  98.02524   94.44136  100.733025 107.8992    95.57442  106.34978\n",
            "  96.10211  112.58493  101.88637   98.96474  101.899574  95.47772\n",
            " 108.00675  116.76012   97.79944   97.47626  117.845985 112.64025\n",
            "  92.89572  100.36632   96.89391  108.944176 106.701096  99.47632\n",
            "  93.87617  112.35951   96.83467  107.74719 ]\n",
            "\n",
            "batch jax default first call\n",
            "CPU times: user 10min 43s, sys: 2.99 s, total: 10min 46s\n",
            "Wall time: 10min 45s\n",
            "batch jax default second call\n",
            "CPU times: user 279 ms, sys: 2 ms, total: 281 ms\n",
            "Wall time: 281 ms\n",
            "[117.34008  105.93037  117.13354   98.280815 104.79451  117.04178\n",
            " 111.9109   107.09849   98.45252   97.8512    96.68642   98.21728\n",
            "  93.4429   117.164894  94.2476    96.47309  106.859344 108.07891\n",
            " 111.64319  104.01584  108.439125  97.67263  106.1951   109.37219\n",
            "  98.36543   93.17829  104.95553  107.21204   93.55455  104.81387\n",
            "  97.9903   117.332054 117.09793   96.438736 104.06102  104.08189\n",
            "  94.45322  118.06995   98.21306   93.453064  98.755936 108.42296\n",
            "  96.125084 106.638275 117.939285  98.15523  117.89131   94.7116\n",
            " 117.48742   96.24068   98.15234  102.07528  109.00425  107.478836\n",
            "  96.7345   116.69099  107.00981  111.89899  109.32806  101.34352\n",
            " 117.27952  102.117386  97.68031   93.76876  101.16061  101.14676\n",
            "  99.0296   108.44193   97.46598  101.93562  104.638794  98.22956\n",
            "  98.02524   94.44136  100.733025 107.8992    95.57442  106.34978\n",
            "  96.10211  112.58493  101.88637   98.96474  101.899574  95.47772\n",
            " 108.00675  116.76012   97.79944   97.47626  117.845985 112.64025\n",
            "  92.89572  100.36632   96.89391  108.944176 106.701096  99.47632\n",
            "  93.87617  112.35951   96.83467  107.74719 ]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7eK9BBU_D4dN",
        "outputId": "8a030420-9b00-4e37-8500-5674afad714f"
      },
      "source": [
        "def batch_sw_v1(unroll=2, NINF=-1e30):\n",
        "    sw_func = sw_v4(unroll=unroll, NINF=NINF)\n",
        "    batch_sw_func = jax.vmap(sw_func, (0, 0, None, None))\n",
        "    return batch_sw_func\n",
        "\n",
        "my_sw_func = batch_sw_v1()\n",
        "print(\"batch jax default first call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "print(\"batch jax default second call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "print(score)\n",
        "print()\n",
        "\n",
        "my_sw_func = jax.jit(batch_sw_v1())\n",
        "print(\"batch jax default first call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "print(\"batch jax default second call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "print(score)\n",
        "print()"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "batch jax default first call\n",
            "CPU times: user 1.04 s, sys: 11 ms, total: 1.05 s\n",
            "Wall time: 1.03 s\n",
            "batch jax default second call\n",
            "CPU times: user 1.04 s, sys: 7.97 ms, total: 1.04 s\n",
            "Wall time: 1.01 s\n",
            "[117.34008  105.93037  117.13354   98.280815 104.79451  117.04178\n",
            " 111.9109   107.09849   98.45252   97.8512    96.68642   98.21728\n",
            "  93.4429   117.164894  94.2476    96.47309  106.859344 108.07891\n",
            " 111.64319  104.01584  108.439125  97.67263  106.1951   109.37219\n",
            "  98.36543   93.17829  104.95553  107.21204   93.55455  104.81387\n",
            "  97.9903   117.332054 117.09793   96.438736 104.06102  104.08189\n",
            "  94.45322  118.06995   98.21306   93.453064  98.755936 108.42296\n",
            "  96.125084 106.638275 117.939285  98.15523  117.89131   94.7116\n",
            " 117.48742   96.24068   98.15234  102.07528  109.00425  107.478836\n",
            "  96.7345   116.69099  107.00981  111.89899  109.32806  101.34352\n",
            " 117.27952  102.117386  97.68031   93.76876  101.16061  101.14676\n",
            "  99.0296   108.44193   97.46598  101.93562  104.638794  98.22956\n",
            "  98.02524   94.44136  100.733025 107.8992    95.57442  106.34978\n",
            "  96.10211  112.58493  101.88637   98.96474  101.899574  95.47772\n",
            " 108.00675  116.76012   97.79944   97.47626  117.845985 112.64025\n",
            "  92.89572  100.36632   96.89391  108.944176 106.701096  99.47632\n",
            "  93.87617  112.35951   96.83467  107.74719 ]\n",
            "\n",
            "batch jax default first call\n",
            "CPU times: user 1.51 s, sys: 10 ms, total: 1.52 s\n",
            "Wall time: 1.5 s\n",
            "batch jax default second call\n",
            "CPU times: user 120 ms, sys: 9 µs, total: 120 ms\n",
            "Wall time: 97 ms\n",
            "[117.34008  105.93037  117.13354   98.280815 104.79451  117.04178\n",
            " 111.9109   107.09849   98.45252   97.8512    96.68642   98.21728\n",
            "  93.4429   117.164894  94.2476    96.47309  106.859344 108.07891\n",
            " 111.64319  104.01584  108.439125  97.67263  106.1951   109.37219\n",
            "  98.36543   93.17829  104.95553  107.21204   93.55455  104.81387\n",
            "  97.9903   117.332054 117.09793   96.438736 104.06102  104.08189\n",
            "  94.45322  118.06995   98.21306   93.453064  98.755936 108.42296\n",
            "  96.125084 106.638275 117.939285  98.15523  117.89131   94.7116\n",
            " 117.48742   96.24068   98.15234  102.07528  109.00425  107.478836\n",
            "  96.7345   116.69099  107.00981  111.89899  109.32806  101.34352\n",
            " 117.27952  102.117386  97.68031   93.76876  101.16061  101.14676\n",
            "  99.0296   108.44193   97.46598  101.93562  104.638794  98.22956\n",
            "  98.02524   94.44136  100.733025 107.8992    95.57442  106.34978\n",
            "  96.10211  112.58493  101.88637   98.96474  101.899574  95.47772\n",
            " 108.00675  116.76012   97.79944   97.47626  117.845985 112.64025\n",
            "  92.89572  100.36632   96.89391  108.944176 106.701096  99.47632\n",
            "  93.87617  112.35951   96.83467  107.74719 ]\n",
            "\n"
          ]
        }
      ]
    }
  ]
}