{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2466d656318d",
   "metadata": {},
   "source": [
    "# ndonnx backend\n",
    "\n",
    "[ndonnx](https://github.com/Quantco/ndonnx) is an array-API-compatible\n",
    "library that builds an ONNX graph symbolically instead of computing\n",
    "eagerly. Use it when you want to build an ONNX model whose computation\n",
    "involves unit-aware quantities.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36fe1da32e3d",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "```bash\n",
    "pip install brainunit[ndonnx]\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17dc4c4ea644",
   "metadata": {},
   "source": [
    "## Graceful import\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "9c9269dcb8fa",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-05-22T03:10:12.602001600Z",
     "start_time": "2026-05-22T03:10:11.459669200Z"
    }
   },
   "source": [
    "import brainunit as u\n",
    "\n",
    "try:\n",
    "    import ndonnx\n",
    "    HAVE_NDONNX = True\n",
    "except ImportError:\n",
    "    HAVE_NDONNX = False\n",
    "    print('ndonnx not installed; install with: pip install brainunit[ndonnx]')\n"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ndonnx not installed; install with: pip install brainunit[ndonnx]\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "id": "8cd717355cbe",
   "metadata": {},
   "source": [
    "## Quick start — symbolic by default\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "68f8f9614e64",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-05-22T03:10:12.623067900Z",
     "start_time": "2026-05-22T03:10:12.613001700Z"
    }
   },
   "source": [
    "if HAVE_NDONNX:\n",
    "    import numpy as np\n",
    "    q = u.Quantity(ndonnx.asarray(np.array([1.0, 2.0, 3.0])), unit=u.meter)\n",
    "    print('backend       =', q.backend)\n",
    "    print('(q + q).backend =', (q + q).backend)\n"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "id": "72302ea06bda",
   "metadata": {},
   "source": [
    "## Unit checking is independent of the backend\n",
    "\n",
    "Dimensional analysis lives on the Python `Quantity`. Mixing units fails the\n",
    "same way regardless of mantissa kind.\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "9cf0c2cbdd57",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-05-22T03:10:12.633920Z",
     "start_time": "2026-05-22T03:10:12.625106400Z"
    }
   },
   "source": [
    "if HAVE_NDONNX:\n",
    "    import numpy as np\n",
    "    from brainunit import UnitMismatchError\n",
    "    a = u.Quantity(ndonnx.asarray(np.array([1.0])), unit=u.meter)\n",
    "    b = u.Quantity(ndonnx.asarray(np.array([1.0])), unit=u.second)\n",
    "    try:\n",
    "        a + b\n",
    "    except UnitMismatchError as exc:\n",
    "        print('expected:', exc)\n"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "id": "a0785c69556d",
   "metadata": {},
   "source": [
    "## Materialization\n",
    "\n",
    "`Quantity.to_numpy()` calls ndonnx's `unwrap_numpy` on the mantissa rather\n",
    "than `np.asarray` (which would give you back a 0-d object wrapper).\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "8a0f4883603e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-05-22T03:10:12.645128300Z",
     "start_time": "2026-05-22T03:10:12.634919200Z"
    }
   },
   "source": [
    "if HAVE_NDONNX:\n",
    "    import numpy as np\n",
    "    q = u.Quantity(ndonnx.asarray(np.array([1.0, 2.0])), unit=u.meter)\n",
    "    eager = q.to_numpy()\n",
    "    print(eager.backend, eager.mantissa)\n"
   ],
   "outputs": [],
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "id": "28658a0b0ea0",
   "metadata": {},
   "source": [
    "## Caveats\n",
    "\n",
    "- ndonnx is still maturing; some `brainunit.math` / `brainunit.linalg`\n",
    "  operations may not have an ndonnx implementation. When that happens, the\n",
    "  ndonnx error surfaces unwrapped — brainunit does not catch it. Consult the\n",
    "  ndonnx documentation for the supported op set.\n",
    "- brainunit does not provide brainunit-level ONNX export helpers. Unit\n",
    "  information lives on the Python `Quantity` object and is not encoded in\n",
    "  the ONNX graph. Export the mantissa using ndonnx's own workflow.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}