jax_vs_map_for_medium.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "37e9c86d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dill\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bea0294b",
   "metadata": {},
   "source": [
    "## Load/create some data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "65d9f3be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "300"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(\"data/pmf_model.dill\", \"rb\") as f:\n",
    "    pmf = dill.load(f)\n",
    "    \n",
    "len(pmf.trace)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aca80cfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" Train-test split \"\"\"\n",
    "train, test = train_test_split(pivoted_ratings.values, frac_test=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eefee9cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(150, 610, 10) (150, 10, 9724)\n"
     ]
    }
   ],
   "source": [
    "U = pmf.trace[\"U\"][-150:]\n",
    "Vt = np.transpose(pmf.trace[\"V\"], axes=[0, 2, 1])[-150:]\n",
    "\n",
    "print(U.shape, Vt.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaeb293b",
   "metadata": {},
   "source": [
    "## Naive Python mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "09d495fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(610, 9724)\n",
      "CPU times: user 8.16 s, sys: 25.2 s, total: 33.4 s\n",
      "Wall time: 53.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "py_map_R = map(np.matmul, U, Vt)\n",
    "py_map_R = np.mean(tuple(py_map_R), axis=0)\n",
    "print(py_map_R.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9436934e",
   "metadata": {},
   "source": [
    "## Failed attempt with `np.tensordot`\n",
    "Jupyter straight up crashed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd6acc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "np_pred_R = np.tensordot(U, Vt, axes = [[2], [1]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd3b672c",
   "metadata": {},
   "source": [
    "## `jax.vmap`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "43414825",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Transform jnp.matmul with jax.vmap, which vectorizes\n",
    "    the function across the axes specified in `in_axes`. \n",
    "    \n",
    "    In other words, `in_axes` specifies the batch dimension\n",
    "    in each function argument --- of which there are two here.\n",
    "\"\"\"\n",
    "\n",
    "batch_matmul = jax.vmap(jnp.matmul, in_axes=(0, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e97b0b9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(610, 9724)\n",
      "CPU times: user 3.46 s, sys: 7.06 s, total: 10.5 s\n",
      "Wall time: 2.91 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "jax_R = batch_matmul(U, Vt)\n",
    "jax_R = jnp.mean(jax_R, axis=0).block_until_ready()\n",
    "\n",
    "print(jax_R.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21f0ebb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "j"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}