Update for recent PyMC: Move from aesara to pytensor

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,15 +16,15 @@
     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
     input_dim = 2
     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

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

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