docs/report/000.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{autolink-concat}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "lambdification",
     "sympy"
    ]
   },
   "source": [
    "::::{margin}\n",
    ":::{card} Square root over arrays with negative values\n",
    "TR-000\n",
    "^^^\n",
    "This notebook investigates how to write a square root function in {mod}`sympy` that computes the positive square root for negative values. The lambdified version of this 'complex square root' should have the same behavior for each computational backend.\n",
    "+++\n",
    "✅&nbsp;[tensorwaves#284](https://github.com/ComPWA/tensorwaves/pull/284)\n",
    ":::\n",
    "::::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Complex square roots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<!-- cspell:disable -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q black==21.5b2 jax==0.2.13 jaxlib==0.1.67 numpy==1.23 sympy==1.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import inspect\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from black import FileMode, format_str\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Negative input values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using {mod}`numpy` as back-end, {mod}`sympy` lambdifies a {func}`~sympy.functions.elementary.miscellaneous.sqrt` to a {obj}`numpy.sqrt`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\sqrt{x}$"
      ],
      "text/plain": [
       "sqrt(x)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = sp.Symbol(\"x\")\n",
    "sqrt_expr = sp.sqrt(x)\n",
    "sqrt_expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return (sqrt(x))\n",
      "\n"
     ]
    }
   ],
   "source": [
    "np_sqrt = sp.lambdify(x, sqrt_expr, \"numpy\")\n",
    "source = inspect.getsource(np_sqrt)\n",
    "print(source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, if input values for the {obj}`numpy.sqrt` are negative, {mod}`numpy` raises a {class}`RuntimeWarning` and returns `NaN`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<lambdifygenerated-1>:2: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return (sqrt(x))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([       nan,        nan, 0.        , 0.70710678, 1.        ])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample = np.linspace(-1, 1, 5)\n",
    "np_sqrt(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want {mod}`numpy` to return imaginary numbers for negative input values, one can use {class}`complex` input data instead (e.g. {doc}`numpy.complex64 <numpy:reference/arrays.scalars>`). Negative values are then treated as lying just above the real axis, so that their square root is a positive imaginary number:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.        +1.j        , 0.        +0.70710677j,\n",
       "       0.        +0.j        , 0.70710677+0.j        ,\n",
       "       1.        +0.j        ], dtype=complex64)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "complex_sample = sample.astype(np.complex64)\n",
    "np_sqrt(complex_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A {func}`sympy.sqrt <sympy.functions.elementary.miscellaneous.sqrt>` lambdified to [JAX](https://jax.rtfd.io) exhibits the same behavior:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return (sqrt(x))\n",
      "\n"
     ]
    }
   ],
   "source": [
    "jax_sqrt = jax.jit(sp.lambdify(x, sqrt_expr, jnp))\n",
    "source = inspect.getsource(jax_sqrt)\n",
    "print(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray([       nan,        nan, 0.        , 0.70710677, 1.        ],            dtype=float32)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax_sqrt(sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-4.3711388e-08+1.j        , -3.0908620e-08+0.70710677j,\n",
       "              0.0000000e+00+0.j        ,  7.0710677e-01+0.j        ,\n",
       "              1.0000000e+00+0.j        ], dtype=complex64)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax_sqrt(complex_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**There is a problem with this approach though**: once input data is complex, _all_ square roots in a larger expression (some amplitude model) compute imaginary solutions for negative values, while this is not always the desired behavior.\n",
    "\n",
    "Take for instance the two square roots appearing in {class}`~ampform.dynamics.phasespace.PhaseSpaceFactor` --- does the $\\sqrt{s}$ also have to be evaluatable for negative $s$?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complex square root"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numpy also offers a special function that evaluates negative values even if the input values are real: {func}`numpy.emath.sqrt`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1j"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.emath.sqrt(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, the {mod}`jax.numpy` API does not interface to {mod}`numpy.emath`. It is possible to decorate {func}`numpy.emath.sqrt` be decorated with {func}`jax.jit`, but that **only works with static, hashable arguments**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception",
     "hide-output",
     "scroll-output"
    ]
   },
   "outputs": [
    {
     "ename": "TracerArrayConversionError",
     "evalue": "The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnfilteredStackTrace\u001b[0m                      Traceback (most recent call last)",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel_launcher.py:17\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaunch_new_instance\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/traitlets/config/application.py:976\u001b[0m, in \u001b[0;36mApplication.launch_instance\u001b[0;34m(cls, argv, **kwargs)\u001b[0m\n\u001b[1;32m    975\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m--> 976\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelapp.py:712\u001b[0m, in \u001b[0;36mIPKernelApp.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    711\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 712\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    713\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/tornado/platform/asyncio.py:199\u001b[0m, in \u001b[0;36mBaseAsyncIOLoop.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    198\u001b[0m     asyncio\u001b[38;5;241m.\u001b[39mset_event_loop(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop)\n\u001b[0;32m--> 199\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43masyncio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_forever\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    200\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mBaseEventLoop.run_forever\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_once\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    571\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36mBaseEventLoop._run_once\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1858\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m         \u001b[43mhandle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/events.py:81\u001b[0m, in \u001b[0;36mHandle._run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_callback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:510\u001b[0m, in \u001b[0;36mKernel.dispatch_queue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    509\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 510\u001b[0m     \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m    511\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:499\u001b[0m, in \u001b[0;36mKernel.process_one\u001b[0;34m(self, wait)\u001b[0m\n\u001b[1;32m    498\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 499\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:406\u001b[0m, in \u001b[0;36mKernel.dispatch_shell\u001b[0;34m(self, msg)\u001b[0m\n\u001b[1;32m    405\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 406\u001b[0m         \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m    407\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:730\u001b[0m, in \u001b[0;36mKernel.execute_request\u001b[0;34m(self, stream, ident, parent)\u001b[0m\n\u001b[1;32m    729\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 730\u001b[0m     reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m    732\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/ipkernel.py:383\u001b[0m, in \u001b[0;36mIPythonKernel.do_execute\u001b[0;34m(self, code, silent, store_history, user_expressions, allow_stdin, cell_id)\u001b[0m\n\u001b[1;32m    382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 383\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43mshell\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    384\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    385\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    386\u001b[0m \u001b[43m        \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    387\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcell_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcell_id\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    388\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    389\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/zmqshell.py:528\u001b[0m, in \u001b[0;36mZMQInteractiveShell.run_cell\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    527\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 528\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:2881\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m   2880\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2881\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2882\u001b[0m \u001b[43m        \u001b[49m\u001b[43mraw_cell\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshell_futures\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\n\u001b[1;32m   2883\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2884\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:2936\u001b[0m, in \u001b[0;36mInteractiveShell._run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m   2935\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2936\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcoro\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2937\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m(coro)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m     \u001b[43mcoro\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3135\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_async\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, transformed_cell, preprocessing_exc_tuple, cell_id)\u001b[0m\n\u001b[1;32m   3133\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3135\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m   3136\u001b[0m        interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m   3138\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3338\u001b[0m, in \u001b[0;36mInteractiveShell.run_ast_nodes\u001b[0;34m(self, nodelist, cell_name, interactivity, compiler, result)\u001b[0m\n\u001b[1;32m   3337\u001b[0m     asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3338\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m   3339\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "Input \u001b[0;32mIn [13]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m jax_csqrt_error \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np\u001b[38;5;241m.\u001b[39memath\u001b[38;5;241m.\u001b[39msqrt, backend\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_csqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/traceback_util.py:143\u001b[0m, in \u001b[0;36mapi_boundary.<locals>.reraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/api.py:426\u001b[0m, in \u001b[0;36m_cpp_jit.<locals>.cache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    425\u001b[0m flat_fun, out_tree \u001b[38;5;241m=\u001b[39m flatten_fun(f, in_tree)\n\u001b[0;32m--> 426\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxla\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxla_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    427\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    428\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    429\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    430\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    431\u001b[0m \u001b[43m    \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mflat_fun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__name__\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    432\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdonated_invars\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    433\u001b[0m out_pytree_def \u001b[38;5;241m=\u001b[39m out_tree()\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1565\u001b[0m, in \u001b[0;36mCallPrimitive.bind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[0;32m-> 1565\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1556\u001b[0m, in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1555\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m maybe_new_sublevel(top_trace):\n\u001b[0;32m-> 1556\u001b[0m   outs \u001b[38;5;241m=\u001b[39m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmap\u001b[39m(full_lower, apply_todos(env_trace_todo(), outs))\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1568\u001b[0m, in \u001b[0;36mCallPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m   1567\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m-> 1568\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:609\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m    608\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 609\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py:578\u001b[0m, in \u001b[0;36m_xla_call_impl\u001b[0;34m(fun, device, backend, name, donated_invars, *args)\u001b[0m\n\u001b[1;32m    577\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_xla_call_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, device, backend, name, donated_invars):\n\u001b[0;32m--> 578\u001b[0m   compiled_fun \u001b[38;5;241m=\u001b[39m \u001b[43m_xla_callable\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    579\u001b[0m \u001b[43m                               \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43munsafe_map\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg_spec\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    580\u001b[0m   \u001b[38;5;28;01mtry\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py:262\u001b[0m, in \u001b[0;36mcache.<locals>.memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 262\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    263\u001b[0m   cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py:652\u001b[0m, in \u001b[0;36m_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m    651\u001b[0m abstract_args, _ \u001b[38;5;241m=\u001b[39m unzip2(arg_specs)\n\u001b[0;32m--> 652\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mabstract_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtransform_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mjit\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(c, core\u001b[38;5;241m.\u001b[39mTracer) \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m consts):\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py:1209\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, transform_name)\u001b[0m\n\u001b[1;32m   1208\u001b[0m main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m ()  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[0;32m-> 1209\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1210\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py:1188\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals)\u001b[0m\n\u001b[1;32m   1187\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[0;32m-> 1188\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1189\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py:166\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    165\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 166\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    167\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    168\u001b[0m   \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m    169\u001b[0m   \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m    170\u001b[0m   \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks only\u001b[39;00m\n\u001b[1;32m    171\u001b[0m   \u001b[38;5;66;03m# after this exception is handled, which can corrupt the global state.\u001b[39;00m\n",
      "File \u001b[0;32m<__array_function__ internals>:180\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py:247\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m    200\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    201\u001b[0m \u001b[38;5;124;03mCompute the square root of x.\u001b[39;00m\n\u001b[1;32m    202\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    245\u001b[0m \u001b[38;5;124;03m-2j\u001b[39;00m\n\u001b[1;32m    246\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 247\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43m_fix_real_lt_zero\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    248\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m nx\u001b[38;5;241m.\u001b[39msqrt(x)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py:134\u001b[0m, in \u001b[0;36m_fix_real_lt_zero\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m    113\u001b[0m \u001b[38;5;124;03m\"\"\"Convert `x` to complex if it has real, negative components.\u001b[39;00m\n\u001b[1;32m    114\u001b[0m \n\u001b[1;32m    115\u001b[0m \u001b[38;5;124;03mOtherwise, output is just the array version of the input (via asarray).\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    132\u001b[0m \n\u001b[1;32m    133\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 134\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43masarray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    135\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(isreal(x) \u001b[38;5;241m&\u001b[39m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m)):\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:472\u001b[0m, in \u001b[0;36mTracer.__array__\u001b[0;34m(self, *args, **kw)\u001b[0m\n\u001b[1;32m    471\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__array__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkw):\n\u001b[0;32m--> 472\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m TracerArrayConversionError(\u001b[38;5;28mself\u001b[39m)\n",
      "\u001b[0;31mUnfilteredStackTrace\u001b[0m: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mTracerArrayConversionError\u001b[0m                Traceback (most recent call last)",
      "Input \u001b[0;32mIn [13]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m jax_csqrt_error \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np\u001b[38;5;241m.\u001b[39memath\u001b[38;5;241m.\u001b[39msqrt, backend\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_csqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m<__array_function__ internals>:180\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py:247\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m    198\u001b[0m \u001b[38;5;129m@array_function_dispatch\u001b[39m(_unary_dispatcher)\n\u001b[1;32m    199\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21msqrt\u001b[39m(x):\n\u001b[1;32m    200\u001b[0m     \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    201\u001b[0m \u001b[38;5;124;03m    Compute the square root of x.\u001b[39;00m\n\u001b[1;32m    202\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    245\u001b[0m \u001b[38;5;124;03m    -2j\u001b[39;00m\n\u001b[1;32m    246\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 247\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[43m_fix_real_lt_zero\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    248\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m nx\u001b[38;5;241m.\u001b[39msqrt(x)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py:134\u001b[0m, in \u001b[0;36m_fix_real_lt_zero\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m    112\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_fix_real_lt_zero\u001b[39m(x):\n\u001b[1;32m    113\u001b[0m     \u001b[38;5;124;03m\"\"\"Convert `x` to complex if it has real, negative components.\u001b[39;00m\n\u001b[1;32m    114\u001b[0m \n\u001b[1;32m    115\u001b[0m \u001b[38;5;124;03m    Otherwise, output is just the array version of the input (via asarray).\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    132\u001b[0m \n\u001b[1;32m    133\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 134\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[43masarray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    135\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(isreal(x) \u001b[38;5;241m&\u001b[39m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m)):\n\u001b[1;32m    136\u001b[0m         x \u001b[38;5;241m=\u001b[39m _tocomplex(x)\n",
      "\u001b[0;31mTracerArrayConversionError\u001b[0m: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)"
     ]
    }
   ],
   "source": [
    "jax_csqrt_error = jax.jit(np.emath.sqrt, backend=\"cpu\")\n",
    "jax_csqrt_error(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.+1.j, dtype=complex64)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax_csqrt = jax.jit(np.emath.sqrt, backend=\"cpu\", static_argnums=0)\n",
    "jax_csqrt(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class 'numpy.ndarray'>, [-1.  -0.5  0.   0.5  1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Input \u001b[0;32mIn [15]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjax_csqrt\u001b[49m\u001b[43m(\u001b[49m\u001b[43msample\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mValueError\u001b[0m: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class 'numpy.ndarray'>, [-1.  -0.5  0.   0.5  1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n"
     ]
    }
   ],
   "source": [
    "jax_csqrt(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Conditional square root"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To be able to control which square roots in the complete expression should be evaluatable for negative values, one could use {class}`~sympy.functions.elementary.piecewise.Piecewise`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{cases} i \\sqrt{- x} & \\text{for}\\: x < 0 \\\\\\sqrt{x} & \\text{otherwise} \\end{cases}$"
      ],
      "text/plain": [
       "Piecewise((I*sqrt(-x), x < 0), (sqrt(x), True))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def complex_sqrt(x: sp.Symbol) -> sp.Expr:\n",
    "    return sp.Piecewise(\n",
    "        (sp.sqrt(-x) * sp.I, x < 0),\n",
    "        (sp.sqrt(x), True),\n",
    "    )\n",
    "\n",
    "\n",
    "complex_sqrt(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 2 i$"
      ],
      "text/plain": [
       "2*I"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 2$"
      ],
      "text/plain": [
       "2"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(\n",
    "    complex_sqrt(-4),\n",
    "    complex_sqrt(+4),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Be careful though when lambdifying this expression: do not use the `__dict__` of the {mod}`numpy` module as backend, but use the module itself instead. When using `__dict__`, {func}`~sympy.utilities.lambdify.lambdify` will return an `if-else` statement, which is inefficient and, worse, will result in problems with {doc}`JAX <jax:index>`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{warning}\n",
    "\n",
    "Do not use the module `__dict__` for the `modules` argument of {func}`~sympy.utilities.lambdify.lambdify`.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return (((1j*sqrt(-x)) if (x < 0) else (sqrt(x))))\n",
      "\n"
     ]
    }
   ],
   "source": [
    "np_complex_sqrt_no_select = sp.lambdify(x, complex_sqrt(x), np.__dict__)\n",
    "source = inspect.getsource(np_complex_sqrt_no_select)\n",
    "print(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1j"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_complex_sqrt_no_select(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "ename": "ConcretizationTypeError",
     "evalue": "Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>\nThe problem arose with the `bool` function. \nWhile tracing the function _lambdifygenerated at <lambdifygenerated-3>:1, transformed by jit., this concrete value was not available in Python because it depends on the value of the arguments to _lambdifygenerated at <lambdifygenerated-3>:1, transformed by jit. at flattened positions [0], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).\n (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnfilteredStackTrace\u001b[0m                      Traceback (most recent call last)",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel_launcher.py:17\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaunch_new_instance\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/traitlets/config/application.py:976\u001b[0m, in \u001b[0;36mApplication.launch_instance\u001b[0;34m(cls, argv, **kwargs)\u001b[0m\n\u001b[1;32m    975\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m--> 976\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelapp.py:712\u001b[0m, in \u001b[0;36mIPKernelApp.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    711\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 712\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    713\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/tornado/platform/asyncio.py:199\u001b[0m, in \u001b[0;36mBaseAsyncIOLoop.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    198\u001b[0m     asyncio\u001b[38;5;241m.\u001b[39mset_event_loop(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop)\n\u001b[0;32m--> 199\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43masyncio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_forever\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    200\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mBaseEventLoop.run_forever\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_once\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    571\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36mBaseEventLoop._run_once\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1858\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m         \u001b[43mhandle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/events.py:81\u001b[0m, in \u001b[0;36mHandle._run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_callback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:510\u001b[0m, in \u001b[0;36mKernel.dispatch_queue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    509\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 510\u001b[0m     \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m    511\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:499\u001b[0m, in \u001b[0;36mKernel.process_one\u001b[0;34m(self, wait)\u001b[0m\n\u001b[1;32m    498\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 499\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:406\u001b[0m, in \u001b[0;36mKernel.dispatch_shell\u001b[0;34m(self, msg)\u001b[0m\n\u001b[1;32m    405\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 406\u001b[0m         \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m    407\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:730\u001b[0m, in \u001b[0;36mKernel.execute_request\u001b[0;34m(self, stream, ident, parent)\u001b[0m\n\u001b[1;32m    729\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 730\u001b[0m     reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m    732\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/ipkernel.py:383\u001b[0m, in \u001b[0;36mIPythonKernel.do_execute\u001b[0;34m(self, code, silent, store_history, user_expressions, allow_stdin, cell_id)\u001b[0m\n\u001b[1;32m    382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 383\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43mshell\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    384\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    385\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    386\u001b[0m \u001b[43m        \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    387\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcell_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcell_id\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    388\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    389\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/zmqshell.py:528\u001b[0m, in \u001b[0;36mZMQInteractiveShell.run_cell\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    527\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 528\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:2881\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m   2880\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2881\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2882\u001b[0m \u001b[43m        \u001b[49m\u001b[43mraw_cell\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshell_futures\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\n\u001b[1;32m   2883\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2884\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:2936\u001b[0m, in \u001b[0;36mInteractiveShell._run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m   2935\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2936\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcoro\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2937\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m(coro)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m     \u001b[43mcoro\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3135\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_async\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, transformed_cell, preprocessing_exc_tuple, cell_id)\u001b[0m\n\u001b[1;32m   3133\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3135\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m   3136\u001b[0m        interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m   3138\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3338\u001b[0m, in \u001b[0;36mInteractiveShell.run_ast_nodes\u001b[0;34m(self, nodelist, cell_name, interactivity, compiler, result)\u001b[0m\n\u001b[1;32m   3337\u001b[0m     asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3338\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m   3339\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "Input \u001b[0;32mIn [20]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m jax_complex_sqrt_no_select \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np_complex_sqrt_no_select)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_complex_sqrt_no_select\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/traceback_util.py:143\u001b[0m, in \u001b[0;36mapi_boundary.<locals>.reraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/api.py:426\u001b[0m, in \u001b[0;36m_cpp_jit.<locals>.cache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    425\u001b[0m flat_fun, out_tree \u001b[38;5;241m=\u001b[39m flatten_fun(f, in_tree)\n\u001b[0;32m--> 426\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxla\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxla_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    427\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    428\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    429\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    430\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    431\u001b[0m \u001b[43m    \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mflat_fun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__name__\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    432\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdonated_invars\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    433\u001b[0m out_pytree_def \u001b[38;5;241m=\u001b[39m out_tree()\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1565\u001b[0m, in \u001b[0;36mCallPrimitive.bind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[0;32m-> 1565\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1556\u001b[0m, in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1555\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m maybe_new_sublevel(top_trace):\n\u001b[0;32m-> 1556\u001b[0m   outs \u001b[38;5;241m=\u001b[39m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmap\u001b[39m(full_lower, apply_todos(env_trace_todo(), outs))\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1568\u001b[0m, in \u001b[0;36mCallPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m   1567\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m-> 1568\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:609\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m    608\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 609\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py:578\u001b[0m, in \u001b[0;36m_xla_call_impl\u001b[0;34m(fun, device, backend, name, donated_invars, *args)\u001b[0m\n\u001b[1;32m    577\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_xla_call_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, device, backend, name, donated_invars):\n\u001b[0;32m--> 578\u001b[0m   compiled_fun \u001b[38;5;241m=\u001b[39m \u001b[43m_xla_callable\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    579\u001b[0m \u001b[43m                               \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43munsafe_map\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg_spec\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    580\u001b[0m   \u001b[38;5;28;01mtry\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py:262\u001b[0m, in \u001b[0;36mcache.<locals>.memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 262\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    263\u001b[0m   cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py:652\u001b[0m, in \u001b[0;36m_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m    651\u001b[0m abstract_args, _ \u001b[38;5;241m=\u001b[39m unzip2(arg_specs)\n\u001b[0;32m--> 652\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mabstract_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtransform_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mjit\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(c, core\u001b[38;5;241m.\u001b[39mTracer) \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m consts):\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py:1209\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, transform_name)\u001b[0m\n\u001b[1;32m   1208\u001b[0m main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m ()  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[0;32m-> 1209\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1210\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py:1188\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals)\u001b[0m\n\u001b[1;32m   1187\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[0;32m-> 1188\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1189\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py:166\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    165\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 166\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    167\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    168\u001b[0m   \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m    169\u001b[0m   \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m    170\u001b[0m   \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks only\u001b[39;00m\n\u001b[1;32m    171\u001b[0m   \u001b[38;5;66;03m# after this exception is handled, which can corrupt the global state.\u001b[39;00m\n",
      "File \u001b[0;32m<lambdifygenerated-3>:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m (((\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx)) \u001b[38;5;28;01mif\u001b[39;00m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m (sqrt(x))))\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:538\u001b[0m, in \u001b[0;36mTracer.__bool__\u001b[0;34m(self)\u001b[0m\n\u001b[0;32m--> 538\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__bool__\u001b[39m(\u001b[38;5;28mself\u001b[39m): \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maval\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_bool\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:960\u001b[0m, in \u001b[0;36mconcretization_function_error.<locals>.error\u001b[0;34m(self, arg)\u001b[0m\n\u001b[1;32m    959\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21merror\u001b[39m(\u001b[38;5;28mself\u001b[39m, arg):\n\u001b[0;32m--> 960\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m ConcretizationTypeError(arg, fname_context)\n",
      "\u001b[0;31mUnfilteredStackTrace\u001b[0m: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>\nThe problem arose with the `bool` function. \nWhile tracing the function _lambdifygenerated at <lambdifygenerated-3>:1, transformed by jit., this concrete value was not available in Python because it depends on the value of the arguments to _lambdifygenerated at <lambdifygenerated-3>:1, transformed by jit. at flattened positions [0], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).\n (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError)\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mConcretizationTypeError\u001b[0m                   Traceback (most recent call last)",
      "Input \u001b[0;32mIn [20]\u001b[0m, in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m jax_complex_sqrt_no_select \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np_complex_sqrt_no_select)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_complex_sqrt_no_select\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m<lambdifygenerated-3>:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m (((\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx)) \u001b[38;5;28;01mif\u001b[39;00m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m (sqrt(x))))\n",
      "\u001b[0;31mConcretizationTypeError\u001b[0m: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>\nThe problem arose with the `bool` function. \nWhile tracing the function _lambdifygenerated at <lambdifygenerated-3>:1, transformed by jit., this concrete value was not available in Python because it depends on the value of the arguments to _lambdifygenerated at <lambdifygenerated-3>:1, transformed by jit. at flattened positions [0], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).\n (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError)"
     ]
    }
   ],
   "source": [
    "jax_complex_sqrt_no_select = jax.jit(np_complex_sqrt_no_select)\n",
    "jax_complex_sqrt_no_select(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When instead using the {mod}`numpy` module (or `\"numpy\"`), {func}`~sympy.utilities.lambdify.lambdify` correctly lambdifies to {func}`numpy.select` to represent the cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np_complex_sqrt = sp.lambdify(x, complex_sqrt(x), np)\n",
    "source = inspect.getsource(np_complex_sqrt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return select(\n",
      "        [less(x, 0), True],\n",
      "        [1j * sqrt(-x), sqrt(x)],\n",
      "        default=nan,\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(format_str(source.replace(\"nan)\", \"nan,)\"), mode=FileMode()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Still, JAX does not handle this correctly. First, lambdifying JAX again results in this `if-else` syntax:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return (((1j*sqrt(-x)) if (x < 0) else (sqrt(x))))\n",
      "\n"
     ]
    }
   ],
   "source": [
    "jnp_complex_sqrt = sp.lambdify(x, complex_sqrt(x), jnp)\n",
    "source = inspect.getsource(jnp_complex_sqrt)\n",
    "print(source)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But even if we lambdify to {mod}`numpy` and decorate the result with a {func}`jax.jit` decorator, the resulting function does not work properly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "jax_complex_sqrt_error = jax.jit(np_complex_sqrt)\n",
    "source = inspect.getsource(jax_complex_sqrt_error)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def _lambdifygenerated(x):\n",
      "    return select(\n",
      "        [less(x, 0), True],\n",
      "        [1j * sqrt(-x), sqrt(x)],\n",
      "        default=nan,\n",
      "    )\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(format_str(source.replace(\"nan)\", \"nan,)\"), mode=FileMode()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception",
     "hide-output"
    ]
   },
   "outputs": [
    {
     "ename": "TracerArrayConversionError",
     "evalue": "The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnfilteredStackTrace\u001b[0m                      Traceback (most recent call last)",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel_launcher.py:17\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaunch_new_instance\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/traitlets/config/application.py:976\u001b[0m, in \u001b[0;36mApplication.launch_instance\u001b[0;34m(cls, argv, **kwargs)\u001b[0m\n\u001b[1;32m    975\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m--> 976\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelapp.py:712\u001b[0m, in \u001b[0;36mIPKernelApp.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    711\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 712\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    713\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/tornado/platform/asyncio.py:199\u001b[0m, in \u001b[0;36mBaseAsyncIOLoop.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    198\u001b[0m     asyncio\u001b[38;5;241m.\u001b[39mset_event_loop(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop)\n\u001b[0;32m--> 199\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43masyncio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_forever\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    200\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mBaseEventLoop.run_forever\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_once\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    571\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36mBaseEventLoop._run_once\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1858\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m         \u001b[43mhandle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/asyncio/events.py:81\u001b[0m, in \u001b[0;36mHandle._run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_callback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:510\u001b[0m, in \u001b[0;36mKernel.dispatch_queue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    509\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 510\u001b[0m     \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m    511\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:499\u001b[0m, in \u001b[0;36mKernel.process_one\u001b[0;34m(self, wait)\u001b[0m\n\u001b[1;32m    498\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 499\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:406\u001b[0m, in \u001b[0;36mKernel.dispatch_shell\u001b[0;34m(self, msg)\u001b[0m\n\u001b[1;32m    405\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 406\u001b[0m         \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m    407\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/kernelbase.py:730\u001b[0m, in \u001b[0;36mKernel.execute_request\u001b[0;34m(self, stream, ident, parent)\u001b[0m\n\u001b[1;32m    729\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 730\u001b[0m     reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m    732\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/ipkernel.py:383\u001b[0m, in \u001b[0;36mIPythonKernel.do_execute\u001b[0;34m(self, code, silent, store_history, user_expressions, allow_stdin, cell_id)\u001b[0m\n\u001b[1;32m    382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 383\u001b[0m     res \u001b[38;5;241m=\u001b[39m \u001b[43mshell\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    384\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    385\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    386\u001b[0m \u001b[43m        \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    387\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcell_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcell_id\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    388\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    389\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/ipykernel/zmqshell.py:528\u001b[0m, in \u001b[0;36mZMQInteractiveShell.run_cell\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    527\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 528\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:2881\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m   2880\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2881\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2882\u001b[0m \u001b[43m        \u001b[49m\u001b[43mraw_cell\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshell_futures\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\n\u001b[1;32m   2883\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2884\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:2936\u001b[0m, in \u001b[0;36mInteractiveShell._run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m   2935\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2936\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcoro\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2937\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m(coro)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m     \u001b[43mcoro\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3135\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_async\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, transformed_cell, preprocessing_exc_tuple, cell_id)\u001b[0m\n\u001b[1;32m   3133\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3135\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m   3136\u001b[0m        interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m   3138\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/IPython/core/interactiveshell.py:3338\u001b[0m, in \u001b[0;36mInteractiveShell.run_ast_nodes\u001b[0;34m(self, nodelist, cell_name, interactivity, compiler, result)\u001b[0m\n\u001b[1;32m   3337\u001b[0m     asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3338\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m   3339\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
      "Input \u001b[0;32mIn [26]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjax_complex_sqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/traceback_util.py:143\u001b[0m, in \u001b[0;36mapi_boundary.<locals>.reraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/api.py:426\u001b[0m, in \u001b[0;36m_cpp_jit.<locals>.cache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    425\u001b[0m flat_fun, out_tree \u001b[38;5;241m=\u001b[39m flatten_fun(f, in_tree)\n\u001b[0;32m--> 426\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxla\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxla_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    427\u001b[0m \u001b[43m    \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    428\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    429\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    430\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    431\u001b[0m \u001b[43m    \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mflat_fun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__name__\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    432\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdonated_invars\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    433\u001b[0m out_pytree_def \u001b[38;5;241m=\u001b[39m out_tree()\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1565\u001b[0m, in \u001b[0;36mCallPrimitive.bind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1564\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[0;32m-> 1565\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1556\u001b[0m, in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m   1555\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m maybe_new_sublevel(top_trace):\n\u001b[0;32m-> 1556\u001b[0m   outs \u001b[38;5;241m=\u001b[39m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1557\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmap\u001b[39m(full_lower, apply_todos(env_trace_todo(), outs))\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:1568\u001b[0m, in \u001b[0;36mCallPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m   1567\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m-> 1568\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:609\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m    608\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 609\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py:578\u001b[0m, in \u001b[0;36m_xla_call_impl\u001b[0;34m(fun, device, backend, name, donated_invars, *args)\u001b[0m\n\u001b[1;32m    577\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_xla_call_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, device, backend, name, donated_invars):\n\u001b[0;32m--> 578\u001b[0m   compiled_fun \u001b[38;5;241m=\u001b[39m \u001b[43m_xla_callable\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    579\u001b[0m \u001b[43m                               \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43munsafe_map\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg_spec\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    580\u001b[0m   \u001b[38;5;28;01mtry\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py:262\u001b[0m, in \u001b[0;36mcache.<locals>.memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m    261\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 262\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    263\u001b[0m   cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py:652\u001b[0m, in \u001b[0;36m_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m    651\u001b[0m abstract_args, _ \u001b[38;5;241m=\u001b[39m unzip2(arg_specs)\n\u001b[0;32m--> 652\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mabstract_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtransform_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mjit\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    653\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(c, core\u001b[38;5;241m.\u001b[39mTracer) \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m consts):\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py:1209\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, transform_name)\u001b[0m\n\u001b[1;32m   1208\u001b[0m main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m ()  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[0;32m-> 1209\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1210\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py:1188\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals)\u001b[0m\n\u001b[1;32m   1187\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[0;32m-> 1188\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1189\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py:166\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    165\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 166\u001b[0m   ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    167\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    168\u001b[0m   \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m    169\u001b[0m   \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m    170\u001b[0m   \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks only\u001b[39;00m\n\u001b[1;32m    171\u001b[0m   \u001b[38;5;66;03m# after this exception is handled, which can corrupt the global state.\u001b[39;00m\n",
      "File \u001b[0;32m<lambdifygenerated-4>:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m (select([\u001b[43mless\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m,\u001b[38;5;28;01mTrue\u001b[39;00m], [\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx),sqrt(x)], default\u001b[38;5;241m=\u001b[39mnan))\n",
      "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py:472\u001b[0m, in \u001b[0;36mTracer.__array__\u001b[0;34m(self, *args, **kw)\u001b[0m\n\u001b[1;32m    471\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__array__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkw):\n\u001b[0;32m--> 472\u001b[0m   \u001b[38;5;28;01mraise\u001b[39;00m TracerArrayConversionError(\u001b[38;5;28mself\u001b[39m)\n",
      "\u001b[0;31mUnfilteredStackTrace\u001b[0m: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mTracerArrayConversionError\u001b[0m                Traceback (most recent call last)",
      "Input \u001b[0;32mIn [26]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjax_complex_sqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m<lambdifygenerated-4>:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m (select([\u001b[43mless\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m,\u001b[38;5;28;01mTrue\u001b[39;00m], [\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx),sqrt(x)], default\u001b[38;5;241m=\u001b[39mnan))\n",
      "\u001b[0;31mTracerArrayConversionError\u001b[0m: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)> (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)"
     ]
    }
   ],
   "source": [
    "jax_complex_sqrt_error(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The very same function in created purely with {mod}`jax.numpy` does work without problems, so it seems this is a SymPy problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def jax_complex_sqrt(x):\n",
    "    return jnp.select(\n",
    "        [jnp.less(x, 0), True],\n",
    "        [1j * jnp.sqrt(-x), jnp.sqrt(x)],\n",
    "        default=jnp.nan,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([0.        +1.j        , 0.        +0.70710677j,\n",
       "             0.        +0.j        , 0.70710677+0.j        ,\n",
       "             1.        +0.j        ], dtype=complex64)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax_complex_sqrt(sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A solution to this is presented in {ref}`report/001:Handle for JAX`."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "toc_visible": true
  },
  "keep_output": true,
  "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.10.13"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}