{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Spox implements the ``inline`` function, which allows inlining existing ONNX models in your own Spox model. Any valid model can be used - it can be produced directly by your code, another converter library, or just prepared beforehand.\n",
    "\n",
    "We'll go over three main applications of ``inline`` - composing existing models, embedding custom operators, and integration with existing converter libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import onnx\n",
    "import onnxruntime\n",
    "from spox import argument, build, inline, Tensor, Var\n",
    "import spox.opset.ai.onnx.v17 as op\n",
    "\n",
    "\n",
    "def run(model: onnx.ModelProto, **kwargs) -> list[np.ndarray]:\n",
    "    return onnxruntime.InferenceSession(model.SerializeToString()).run(\n",
    "        None,\n",
    "        {k: np.array(v) for k, v in kwargs.items()}\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Composition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, prepare some example models to compose in Spox. We'll build them into ONNX, moving them out of Spox entirely.\n",
    "The same principle can be applied to existing models for which you have access to the ``onnx.ModelProto``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x, y = argument(Tensor(float, (None,))), argument(Tensor(float, (None,)))\n",
    "z = op.add(op.div(x, y), op.div(y, x))  # x/y + y/x\n",
    "harmonic_model = build({'x': x, 'y': y}, {'z': z})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "v, t = argument(Tensor(float, (None,))), argument(Tensor(float, ()))\n",
    "w = op.abs(op.sub(v, t))  # |v - t|\n",
    "dist_model = build({'v': v, 't': t}, {'w': w})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compose the models we'll create a new set of arguments and pass them into the functions returned by `inline(harmonic_model)` and `inline(dist_model)`. Those can be either passed positionally or as keyword arguments, following the naming in the model input list. A dictionary is returned, with ordering as in the model output list and keys being the output names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "a, b = argument(Tensor(float, ('N',))), argument(Tensor(float, ('N',)))\n",
    "harmonic_res = inline(harmonic_model)(x=a, y=b)  # Use kwargs & dict result\n",
    "assert list(harmonic_res.keys()) == ['z']\n",
    "c = harmonic_res['z']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "(d,) = inline(dist_model)(c, op.constant(value=np.array(2.0))).values()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In summary, using the models $z = \\frac{x}{y} + \\frac{y}{x}$ and $w = v - t$ we constructed $d = \\left| \\frac{a}{b} + \\frac{b}{a} - 2 \\right|$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "harmonic_dist2_model = build({'a': a, 'b': b}, {'d': d})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom operators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not all operators may be pre-generated in Spox. A simple way to embed a custom operator in Spox is to use ``inline`` with a hand-crafted call with the operator. A wrapper can also perform basic type checking and process some attributes. We'll use the example of the non-standard `Inverse` operator, available in `onnxruntime`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def inverse(m: Var) -> Var:\n",
    "    # Asserts argument is a tensor and checks its dtype\n",
    "    dtype = m.unwrap_tensor().dtype\n",
    "    # Do a basic type check\n",
    "    if dtype not in (np.float32, np.float64):\n",
    "        raise TypeError(\"Input element type expected to be float32/float64.\")\n",
    "    # We'll use this to construct type information for inlined model input/output\n",
    "    tensor_dtype = onnx.helper.np_dtype_to_tensor_dtype(dtype)\n",
    "    # Construct the inlined model\n",
    "    inverse_model = onnx.helper.make_model(\n",
    "        onnx.helper.make_graph(\n",
    "            [onnx.helper.make_node(\"Inverse\", [\"i\"], [\"o\"], domain=\"com.microsoft\")],\n",
    "            \"inverse_graph\",\n",
    "            [onnx.helper.make_tensor_value_info(\"i\", tensor_dtype, (None, None))],\n",
    "            [onnx.helper.make_tensor_value_info(\"o\", tensor_dtype, (None, None))]\n",
    "        ),\n",
    "        opset_imports=[\n",
    "            onnx.helper.make_opsetid(\"\", 16),\n",
    "            onnx.helper.make_opsetid(\"com.microsoft\", 1)\n",
    "        ]\n",
    "    )\n",
    "    # Inline the model into Spox - unpack the only output into m1\n",
    "    (m1,) = inline(inverse_model)(m).values()\n",
    "    return m1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, Spox will use the type information in the constructed model to check that the passed argument is a matrix (tensor of rank 2). We additionally assert the dtype is `float32` or `float64`.\n",
    "\n",
    "We can now use the wrapper in Spox."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "f = argument(Tensor(float, ('N', 'N')))\n",
    "g = op.matmul(f, inverse(f))\n",
    "eye_model = build({'f': f}, {'g': g})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model performs the toy computation of $F F^{-1}$, which obviously should equal $\\mathbf{I}$ (unless $F$ isn't invertible).\n",
    "We can check this is the case when ``eye_model`` is run in ONNX Runtime:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1.00000000e+00, 6.66133815e-16],\n",
       "       [0.00000000e+00, 1.00000000e+00]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(result,) = run(eye_model, f=np.array([[3., 7.], [1., 3.]]))\n",
    "assert np.isclose(result, np.eye(2)).all()\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converter libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, it is worth noting how Spox can use the product of other converter libraries. In practice this reduces to a simple routine: it is enough to make the existing library produce an `onnx.ModelProto` which can be passed to `inline`. The only requirement is knowledge of the signature (model inputs/outputs). Some metadata written to the model may be lost (like e.g. docstrings).\n",
    "\n",
    "For example, we could use `skl_model = skl2onnx.to_onnx(Pipeline([...]))` and then pass this to Spox via `inline(skl_model)(...)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}