Remove Theano support in favor of PyTensor

pymc-devs · Dec 14, 2022 · fa067e8 · fa067e8
1 parent d9a3b2c
commit fa067e8
Show file tree

Hide file tree

Showing 11 changed files with 85 additions and 91 deletions.
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -28,7 +28,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install Dependences
       run: |
-        conda install --yes conda-build conda-verify pytest pytest-cov hypothesis statsmodels aesara c-compiler
+        conda install --yes conda-build conda-verify pytest pytest-cov hypothesis statsmodels pytensor c-compiler
     - name: Build package
       run: |
         conda build --variants "{python: [${{ matrix.python-version }}]}" ./sunode/conda

diff --git a/README.md b/README.md
@@ -13,7 +13,7 @@ functions using symbolic differentiation and common subexpression elimination.
 In either case the functions are compiled using numba and the resulting
 C-function is passed to sunode, so that there is no python overhead.
 
-`sunode` comes with an Aesara wrapper so that parameters of an ODE can be estimated
+`sunode` comes with an PyTensor wrapper so that parameters of an ODE can be estimated
 using PyMC.
 
 ### Installation
@@ -134,7 +134,7 @@ We'll use some time artificial data:
 ```python
 import numpy as np
 import sunode
-import sunode.wrappers.as_aesara
+import sunode.wrappers.as_pytensor
 import pymc as pm
 
 times = np.arange(1900,1921,1)
@@ -186,17 +186,17 @@ with pm.Model() as model:
     gamma = pm.Deterministic('gamma', freq / speed_ratio / ratio)
     delta = pm.Deterministic('delta', freq / speed_ratio / fixed_hares / ratio)
 
-    y_hat, _, problem, solver, _, _ = sunode.wrappers.as_aesara.solve_ivp(
+    y_hat, _, problem, solver, _, _ = sunode.wrappers.as_pytensor.solve_ivp(
         y0={
 	    # The initial conditions of the ode. Each variable
-	    # needs to specify a theano or numpy variable and a shape.
+	    # needs to specify a PyTensor or numpy variable and a shape.
 	    # This dict can be nested.
             'hares': (hares_start, ()),
             'lynx': (lynx_start, ()),
         },
         params={
 	    # Each parameter of the ode. sunode will only compute derivatives
-	    # with respect to theano variables. The shape needs to be specified
+	    # with respect to PyTensor variables. The shape needs to be specified
 	    # as well. It it infered automatically for numpy variables.
 	    # This dict can be nested.
             'alpha': (alpha, ()),

diff --git a/conda/meta.yaml b/conda/meta.yaml
@@ -42,7 +42,7 @@ test:
     - pytest
     - hypothesis
     - statsmodels
-    - aesara
+    - pytensor
   commands:
     - pytest --pyargs sunode
 

diff --git a/doc/source/quickstart_pymc.rst b/doc/source/quickstart_pymc.rst
@@ -42,7 +42,7 @@ We'll use some time artificial data:::
 
     import numpy as np
     import sunode
-    import sunode.wrappers.as_aesara
+    import sunode.wrappers.as_pytensor
     import pymc as pm
 
     times = np.arange(1900,1921,1)
@@ -124,7 +124,7 @@ Now, we define the names, (symbolic) values and shapes of the parameters and ini
 We solve the ODE using the ``solve_ivp`` function from sunode::
 
     with model:
-        from sunode.wrappers.as_aesara import solve_ivp
+        from sunode.wrappers.as_pytensor import solve_ivp
         solution, *_ = solve_ivp(
             y0=y0,
             params=params,

diff --git a/doc/source/without_pymc.rst b/doc/source/without_pymc.rst
@@ -12,7 +12,7 @@ current number, and die when eaten by a lynx. We get:
 .. math::
    \frac{dH}{dt} = \alpha H - \beta LH \\ \frac{dL}{dt} = \delta LH - \gamma L
 
-If we want to solve this ODE without the support of Aesara or PyMC, we need to
+If we want to solve this ODE without the support of PyTensor or PyMC, we need to
 first declare the parameters and states we are using. We have four parameters
 and two states, and each one is a scalar values, so it has shape ()::
 
@@ -71,7 +71,7 @@ ODE might look like this::
    This right-hand-side function is usually only called once to collect the
    sympy expressions of the derivatives. Control flow within this function
    might behave in unexpected ways if you are new to this concept. It is the
-   same thing as with theano, pytorch or tensorflow in graph mode. This means
+   same thing as with PyTensor, pytorch or tensorflow in graph mode. This means
    that something like this will **not** work as expected::
 
        value = 1

diff --git a/mypy.ini b/mypy.ini
@@ -20,10 +20,10 @@ ignore_missing_imports = True
 [mypy-pandas]
 ignore_missing_imports = True
 
-[mypy-theano]
+[mypy-pytensor]
 ignore_missing_imports = True
 
-[mypy-theano.tensor]
+[mypy-pytensor.tensor]
 ignore_missing_imports = True
 
 [mypy-sympy.printing.pycode]

diff --git a/notebooks/from_sympy.ipynb b/notebooks/from_sympy.ipynb
@@ -22,17 +22,18 @@
     "\n",
     "import sunode.symode.paramset\n",
     "import sunode.symode.problem\n",
-    "import sunode.wrappers.as_aesara\n",
+    "import sunode.wrappers.as_pytensor\n",
     "\n",
-    "import aesara\n",
-    "import aesara.tensor as aet"
+    "import pytensor\n",
+    "import pytensor.tensor as pt"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Define ode using sympy and aesara"
+    "### Define ODE using sympy and Pytensor"
    ]
   },
   {
@@ -41,8 +42,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "b = aet.dvector('b')\n",
-    "d = aet.dvector('d')\n",
+    "b = pt.dvector('b')\n",
+    "d = pt.dvector('d')\n",
     "\n",
     "def rhs(t, y, params):\n",
     "    return {\n",
@@ -53,7 +54,7 @@
     "    }\n",
     "\n",
     "\n",
-    "solution, problem, solver = sunode.wrappers.as_aesara.solve_ivp(\n",
+    "solution, problem, solver = sunode.wrappers.as_pytensor.solve_ivp(\n",
     "    t0=0,\n",
     "    y0={\n",
     "        'a': (np.arange(3, dtype=float) + d[0] ** 2, 3),\n",
@@ -185,7 +186,7 @@
    "outputs": [],
    "source": [
     "val = (solution ** 2).sum()\n",
-    "grads = aet.grad(val, [b, d])"
+    "grads = pt.grad(val, [b, d])"
    ]
   },
   {
@@ -194,8 +195,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "func = aesara.function([b, d], [val] + grads)\n",
-    "func2 = aesara.function([b, d], [val])"
+    "func = pytensor.function([b, d], [val] + grads)\n",
+    "func2 = pytensor.function([b, d], [val])"
    ]
   },
   {
@@ -661,10 +662,11 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Integrate into aesara and check gradients"
+    "### Integrate into PyTensor and check gradients"
    ]
   },
   {
@@ -682,23 +684,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "params = aet.dvector('params')\n",
-    "y0 = aet.dvector('y0')\n",
-    "solve_ode = sunode.wrappers.as_aesara.SolveODEAdjoint(solver, 0, tvals)\n",
+    "params = pt.dvector('params')\n",
+    "y0 = pt.dvector('y0')\n",
+    "solve_ode = sunode.wrappers.as_pytensor.SolveODEAdjoint(solver, 0, tvals)\n",
     "solution = solve_ode(y0, params)\n",
     "\n",
     "loss = (solution ** 2).sum()\n",
-    "grad_p, grad_y0 = aet.grad(loss, [params, y0])\n",
-    "func = aesara.function([y0, params], [loss, grad_p, grad_y0])\n",
+    "grad_p, grad_y0 = pt.grad(loss, [params, y0])\n",
+    "func = pytensor.function([y0, params], [loss, grad_p, grad_y0])\n",
     "\n",
     "# Explicit solution\n",
     "loss = ((\n",
     "    ((0.5 * tvals ** 2 * params[1] + tvals * y0[1]) + y0[0]) ** 2\n",
     "    + (tvals * params[1] + y0[1]) ** 2\n",
     ")).sum()\n",
-    "grad_p, grad_y0 = aet.grad(loss, [params, y0])\n",
+    "grad_p, grad_y0 = pt.grad(loss, [params, y0])\n",
     "\n",
-    "func2 = aesara.function([y0, params], [loss, grad_p, grad_y0])"
+    "func2 = pytensor.function([y0, params], [loss, grad_p, grad_y0])"
    ]
   },
   {
@@ -842,7 +844,7 @@
     "    params = pm.Normal('params', sigma=10, shape=ode.n_params)\n",
     "    y0 = pm.Normal('y0', shape=ode.n_states)\n",
     "    \n",
-    "    solve_ode = sunode.wrappers.as_aesara.SolveODEAdjoint(solver, 0, tvals)\n",
+    "    solve_ode = sunode.wrappers.as_pytensor.SolveODEAdjoint(solver, 0, tvals)\n",
     "    mu = solve_ode(y0, params)\n",
     "    error = 0.8 * np.random.randn(len(tvals))\n",
     "    pm.Normal('y', mu=mu[:, 0], sigma=0.8, observed=tvals ** 2 + tvals + 5 + error)\n",

diff --git a/notebooks/pymc_model.ipynb b/notebooks/pymc_model.ipynb
@@ -25,7 +25,7 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "os.environ[\"AESARA_FLAGS\"] = \"floatX=float64\""
+    "os.environ[\"pytensor_FLAGS\"] = \"floatX=float64\""
    ]
   },
   {
@@ -48,8 +48,8 @@
     "from scipy.integrate import odeint\n",
     "\n",
     "import pymc as pm\n",
-    "import aesara\n",
-    "import aesara.tensor as aet\n",
+    "import pytensor\n",
+    "import pytensor.tensor as pt\n",
     "\n",
     "# this notebook show DEBUG log messages\n",
     "# logging.getLogger('pymc').setLevel(logging.DEBUG)\n",
@@ -150,9 +150,9 @@
    "outputs": [],
    "source": [
     "# To demonstrate that test-value computation works, but also for debugging\n",
-    "aesara.config.compute_test_value = 'raise'\n",
-    "aesara.config.exception_verbosity = 'high'\n",
-    "aesara.config.traceback.limit = 5"
+    "pytensor.config.compute_test_value = 'raise'\n",
+    "pytensor.config.exception_verbosity = 'high'\n",
+    "pytensor.config.traceback.limit = 5"
    ]
   },
   {
@@ -176,7 +176,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import sunode.wrappers.as_aesara\n",
+    "import sunode.wrappers.as_pytensor\n",
     "\n",
     "\n",
     "def get_model_sunode():\n",
@@ -187,7 +187,7 @@
     "        s0 = pm.Normal('red_0', mu=10, sigma=2)\n",
     "        extra = pm.Normal('extra', shape=n_extra)\n",
     "\n",
-    "        y_hat, _, _ = sunode.wrappers.as_aesara.solve_ivp(\n",
+    "        y_hat, _, _ = sunode.wrappers.as_pytensor.solve_ivp(\n",
     "            y0={\n",
     "                'S': (s0, ()),  # TODO Infer shape from model?\n",
     "                'P': np.array(y0_true[1], dtype='d'),\n",
@@ -196,7 +196,7 @@
     "            params={\n",
     "                'K_S': (K_S, ()),\n",
     "                'vmax': (vmax, ()),\n",
-    "                'tmp': np.zeros(1),  # TODO aesara wants at least one fixed param\n",
+    "                'tmp': np.zeros(1),  # TODO pytensor wants at least one fixed param\n",
     "                'extra_p': (extra, (n_extra,))\n",
     "            },\n",
     "            rhs=reaction_sympy,\n",
@@ -235,7 +235,7 @@
     "        \n",
     "        extra = pm.Normal('extra', shape=n_extra)\n",
     "\n",
-    "        y_hat, problem, _ = sunode.wrappers.as_aesara.solve_ivp(\n",
+    "        y_hat, problem, _ = sunode.wrappers.as_pytensor.solve_ivp(\n",
     "            y0={\n",
     "                'S': (s0, ()),  # TODO Infer shape from model?\n",
     "                'P': np.array(y0_true[1], dtype='d'),\n",
@@ -244,7 +244,7 @@
     "            params={\n",
     "                'K_S': (K_S, ()),\n",
     "                'vmax': (vmax, ()),\n",
-    "                'tmp': np.zeros(1),  # TODO aesara wants at least one fixed param\n",
+    "                'tmp': np.zeros(1),  # TODO pytensor wants at least one fixed param\n",
     "                'extra_p': (extra, (n_extra,))\n",
     "            },\n",
     "            rhs=reaction_sympy,\n",
@@ -305,22 +305,22 @@
     "    \n",
     "    # create a test function for evaluating the logp value\n",
     "    print('Compiling f_logpt')\n",
-    "    f_logpt = aesara.function(\n",
+    "    f_logpt = pytensor.function(\n",
     "        inputs=t_inputs,\n",
     "        outputs=[pmodel.logpt],\n",
     "        # with float32, allow downcast because the forward integration is always float64\n",
-    "        allow_input_downcast=(aesara.config.floatX == 'float32')\n",
+    "        allow_input_downcast=(pytensor.config.floatX == 'float32')\n",
     "    )\n",
     "    print(f'Test logpt:')\n",
     "    print(f_logpt(*test_inputs))\n",
     "    \n",
     "    # and another test function for evaluating the gradient\n",
     "    print('Compiling f_logpt')\n",
-    "    f_grad = aesara.function(\n",
+    "    f_grad = pytensor.function(\n",
     "        inputs=t_inputs,\n",
-    "        outputs=aet.grad(pmodel.logpt, t_inputs),\n",
+    "        outputs=pt.grad(pmodel.logpt, t_inputs),\n",
     "        # with float32, allow downcast because the forward integration is always float64\n",
-    "        allow_input_downcast=(aesara.config.floatX == 'float32')\n",
+    "        allow_input_downcast=(pytensor.config.floatX == 'float32')\n",
     "    )\n",
     "    print(f'Test gradient:')\n",
     "    print(f_grad(*test_inputs))\n",
@@ -984,7 +984,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "aesara.printing.pydotprint(aet.grad(model_sunode.logpt, model_sunode.vmax), 'ODE_API_shapes_and_benchmarking.png')\n",
+    "pytensor.printing.pydotprint(pt.grad(model_sunode.logpt, model_sunode.vmax), 'ODE_API_shapes_and_benchmarking.png')\n",
     "IPython.display.Image('ODE_API_shapes_and_benchmarking.png')"
    ]
   },
@@ -1007,7 +1007,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from aesara import d3viz\n",
+    "from pytensor import d3viz\n",
     "d3viz.d3viz(model.logpt, 'ODE_API_shapes_and_benchmarking.html')"
    ]
   },
@@ -1041,4 +1041,3 @@
  "nbformat": 4,
  "nbformat_minor": 4
 }
-
diff --git a/sunode/test_aesara.py → sunode/test_pytensor.py b/sunode/test_aesara.py → sunode/test_pytensor.py
@@ -1,6 +1,6 @@
 import numpy as np
-import aesara
-import aesara.tensor as aet
+import pytensor
+import pytensor.tensor as pt
 
 import sunode.wrappers
 
@@ -12,7 +12,7 @@ def dydt_dict(t, y, p):
             'B': y.B,
             'C': y.C,
         }
-    A = aet.dscalar("A")
+    A = pt.dscalar("A")
     A.tag.test_value = np.array(0.9)
 
 
@@ -30,7 +30,7 @@ def dydt_dict(t, y, p):
         'extra': np.array([0.])
     }
 
-    solution, *_ = sunode.wrappers.as_aesara.solve_ivp(
+    solution, *_ = sunode.wrappers.as_pytensor.solve_ivp(
         y0=y0,
         params=params,
         rhs=dydt_dict,
@@ -40,5 +40,5 @@ def dydt_dict(t, y, p):
         solver_kwargs=dict(sens_mode="simultaneous")
     )
 
-    func = aesara.function([A], [solution["A"], solution["B"]])
+    func = pytensor.function([A], [solution["A"], solution["B"]])
     assert func(0.2)[0].shape == time.shape
diff --git a/sunode/wrappers/__init__.py b/sunode/wrappers/__init__.py
@@ -1,4 +1,4 @@
-from . import as_aesara
-from . import as_aesara as as_theano
+from . import as_pytensor
+from . import as_pytensor as as_aesara
 
-__all__ = ['as_theano', 'as_aesara']
+__all__ = ['as_aesara', 'as_pytensor']