notebooks/02_basic_dev_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inside `ulinalg`\n",
    "Inside `ulinalg`, you would define a method consisting of an argument replacer and argument extractor.\n",
    "\n",
    "- The argument extractor is the simpler of the two: It \"extracts\" the array arguments from the method as a `tuple`.\n",
    "- The argument replacer \"replaces\" the array arguments inside args, kwargs with the supplied arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uarray import create_multimethod, all_of_type\n",
    "import unumpy as unp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_argreplacer(args, kwargs, arrays):\n",
    "    out_args = arrays + args[2:]\n",
    "    return out_args, kwargs\n",
    "\n",
    "@create_multimethod(solve_argreplacer)\n",
    "@all_of_type(unp.ndarray)\n",
    "def solve(a, b, sym_pos=False, lower=False, overwrite_a=False, overwrite_b=False, debug=None, check_finite=True, assume_a='gen', transposed=False):\n",
    "    return (a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inside a Numpy backend for `ulinalg` (ideally, `scipy.linalg` itself)\n",
    "Here, you register the implementation for the backend itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uarray import register_implementation\n",
    "from unumpy.numpy_backend import NumpyBackend, compat_check as np_cc\n",
    "import scipy.linalg as linalg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function scipy.linalg.basic.solve(a, b, sym_pos=False, lower=False, overwrite_a=False, overwrite_b=False, debug=None, check_finite=True, assume_a='gen', transposed=False)>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "register_implementation(solve, NumpyBackend, compat_check=np_cc)(linalg.solve)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inside PyTorch or a PyTorch backend for `ulinalg` (ideally, PyTorch itself)\n",
    "Here, you have to perform some translation, because the PyTorch API isn't 1:1 with the NumPy API. See [this documentation page](https://pytorch.org/docs/stable/torch.html#torch.gesv). In practice, you would select the function based on the input arguments, but let's ignore that for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uarray import register_implementation\n",
    "from unumpy.torch_backend import TorchBackend, compat_check as torch_cc\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_impl(a, b, sym_pos=False, lower=False, overwrite_a=False, overwrite_b=False, debug=None, check_finite=True, assume_a='gen', transposed=False):\n",
    "    return torch.gesv(b, a)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.solve_impl(a, b, sym_pos=False, lower=False, overwrite_a=False, overwrite_b=False, debug=None, check_finite=True, assume_a='gen', transposed=False)>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "register_implementation(solve, TorchBackend, compat_check=torch_cc)(solve_impl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# User Experience\n",
    "The user simply imports the right library and uses everything as normal. Let's try solving this fruit puzzle:\n",
    "\n",
    "<img src=\"fruit-puzzle.jpg\" width=400>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = [[3.0, 0.0, 0.0], [1.0, 2.0, 0.0], [0.0, 1.0, -2.0]]\n",
    "b = [[30.0], [18.0], [2.0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "ta = torch.tensor(a, dtype=torch.float32); tb = torch.tensor(b, dtype=torch.float32)\n",
    "\n",
    "tx = solve(ta, tb)\n",
    "tsol = tx[0] + tx[1] + tx[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([15.])\n",
      "<class 'torch.Tensor'>\n"
     ]
    }
   ],
   "source": [
    "print(tsol)\n",
    "print(type(tsol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "na = np.array(a); nb = np.array(b)\n",
    "\n",
    "nx = solve(na, nb)\n",
    "nsol = nx[0] + nx[1] + nx[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[15.]\n",
      "<class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "print(nsol)\n",
    "print(type(nsol))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manually choosing the back-end\n",
    "It's also possible to use context managers to safely set the back-end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'numpy.ndarray'>\n",
      "<class 'torch.Tensor'>\n"
     ]
    }
   ],
   "source": [
    "import uarray as ua\n",
    "with ua.set_backend(NumpyBackend, coerce=True):\n",
    "    print(type(solve(a, b)))\n",
    "        \n",
    "with ua.set_backend(TorchBackend, coerce=True):\n",
    "    print(type(solve(a, b)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}