From 6b897de4bdc537628a608727ee62ea592c830294 Mon Sep 17 00:00:00 2001
From: Linus Seelinger <mail@linusseelinger.de>
Date: Thu, 27 Apr 2023 21:49:26 +0200
Subject: [PATCH] Update for recent PyMC: Move from aesara to pytensor

---
 clients/python/pymc-client.py |  8 ++++----
 setup.py                      |  3 ---
 umbridge/pymc.py              | 16 ++++++++--------
 3 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/clients/python/pymc-client.py b/clients/python/pymc-client.py
index 01371bb..10d61f4 100644
--- a/clients/python/pymc-client.py
+++ b/clients/python/pymc-client.py
@@ -1,8 +1,8 @@
 import argparse
 import numpy as np
 import pymc as pm
-import aesara.tensor as at
-from aesara.gradient import verify_grad
+from pytensor import tensor as pt
+from pytensor.gradient import verify_grad
 import arviz as az
 import matplotlib.pyplot as plt
 from umbridge.pymc import UmbridgeOp
@@ -16,7 +16,7 @@
     args = parser.parse_args()
     print(f"Connecting to host URL {args.url}")
 
-    # Set up an aesara op connecting to UM-Bridge model
+    # Set up an pytensor op connecting to UM-Bridge model
     op = UmbridgeOp(args.url, "posterior")
 
     # Define input parameter
@@ -24,7 +24,7 @@
     input_val = [0.0, 10.0]
 
     # Evaluate model with input parameter
-    op_application = op(at.as_tensor_variable(input_val))
+    op_application = op(pt.as_tensor_variable(input_val))
     print(f"Model output: {op_application.eval()}")
 
     # Verify gradient
diff --git a/setup.py b/setup.py
index 1ed381b..9694804 100644
--- a/setup.py
+++ b/setup.py
@@ -12,9 +12,6 @@
     long_description=long_description,
     long_description_content_type="text/markdown",
     install_requires=["aiohttp", "requests", "asyncio"],
-    extras_require = {
-        'pymc':  ["aesara"]
-    },
     url="https://github.com/UM-Bridge/umbridge",
     packages=setuptools.find_packages(),
     classifiers=(                                 # Classifiers help people find your
diff --git a/umbridge/pymc.py b/umbridge/pymc.py
index 7294704..4ca108e 100644
--- a/umbridge/pymc.py
+++ b/umbridge/pymc.py
@@ -1,12 +1,12 @@
-import aesara.tensor as at
+from pytensor import tensor as pt
 import umbridge
 import numpy as np
 
 
-class UmbridgeGradOp(at.Op):
+class UmbridgeGradOp(pt.Op):
 
-    itypes = [at.dvector, at.dvector]
-    otypes = [at.dvector]
+    itypes = [pt.dvector, pt.dvector]
+    otypes = [pt.dvector]
 
     def __init__(self, umbridge_model, config):
         self.umbridge_model = umbridge_model
@@ -16,17 +16,17 @@ def perform(self, node, inputs_var, output_storage):
         grad = self.umbridge_model.gradient(0, 0, [inputs_var[0].tolist()], inputs_var[1].tolist(), self.config)
         output_storage[0][0] = np.asarray(grad).astype('float64')
 
-class UmbridgeOp(at.Op):
+class UmbridgeOp(pt.Op):
 
-    itypes = [at.dvector]
-    otypes = [at.dvector]
+    itypes = [pt.dvector]
+    otypes = [pt.dvector]
 
     # Take model URL in constructor
     def __init__(self, url, name, config = {}):
         self.umbridge_model = umbridge.HTTPModel(url, name)
         self.config = config
         # For now, make sure model takes a single input vector and returns a single output vector.
-        # More could be supported, but needs improved aesara op.
+        # More could be supported, but needs improved pytensor op.
         # (i.e. adjust input/output types according to UM-Bridge model, pass through multiple vectors etc.)
         assert len(self.umbridge_model.get_input_sizes(config)) == 1
         assert len(self.umbridge_model.get_output_sizes(config)) == 1