PyMC 4 & Aesara Cannot convert Elemwise{pow,no_inplace}.0 to a tensor variable

Hi Everyone. I’ve been struggling for a few days trying to migrate a script that previously ran on PyMC3 + Theano to a version that works with PyMC4 + Aesara - hoping to take advantage of GPU acceleration.

The snippet below isolates my current issue

import numpy as np
import aesara.tensor as at
import pymc as pm
from pymc import Beta

def geometric_adstock_at(x, alpha=0,L=30, normalize=True):
    w = at.power(alpha, np.arange(L))
    xx = at.stack([at.concatenate([at.zeros(i), x[:x.shape[0]-i]]) for i in range(L)])
    
    if not normalize:
        y = at.dot(w,xx)
    else:
        y = at.dot(w/at.sum(w),xx)
    return y


x_test = np.ones(100)*10000

with pm.Model() as test_model:
    alpha_test = Beta('alpha_test',3,3)
    geometric_adstock_at(x_test, alpha=alpha_test)

This results in

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[1], line 21
     19 with pm.Model() as test_model:
     20     alpha_test = Beta('alpha_test',3,3)
---> 21     geometric_adstock_at(x_test, alpha=alpha_test)

Cell In[1], line 13, in geometric_adstock_at(x, alpha, L, normalize)
     11     y = at.dot(w,xx)
     12 else:
---> 13     y = at.dot(w/at.sum(w),xx)
     14 return y

File ~/SageMaker/custom_39_venv/lib/python3.9/site-packages/aesara/tensor/math.py:2507, in sum(input, axis, dtype, keepdims, acc_dtype)
   2487 def sum(input, axis=None, dtype=None, keepdims=False, acc_dtype=None):
   2488     """
   2489     Computes the sum along the given axis(es) of a tensor `input`.
   2490 
   (...)
   2504 
   2505     """
-> 2507     out = Sum(axis=axis, dtype=dtype, acc_dtype=acc_dtype)(input)
   2509     if keepdims:
   2510         out = makeKeepDims(input, out, axis)

File ~/SageMaker/custom_39_venv/lib/python3.9/site-packages/aesara/graph/op.py:297, in Op.__call__(self, *inputs, **kwargs)
    255 r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.
    256 
    257 This method is just a wrapper around :meth:`Op.make_node`.
   (...)
    294 
    295 """
    296 return_list = kwargs.pop("return_list", False)
--> 297 node = self.make_node(*inputs, **kwargs)
    299 if config.compute_test_value != "off":
    300     compute_test_value(node)

File ~/SageMaker/custom_39_venv/lib/python3.9/site-packages/aesara/tensor/elemwise.py:1413, in CAReduce.make_node(self, input)
   1412 def make_node(self, input):
-> 1413     input = as_tensor_variable(input)
   1414     inp_dims = input.type.ndim
   1415     inp_dtype = input.type.dtype

File ~/SageMaker/custom_39_venv/lib/python3.9/site-packages/aesara/tensor/__init__.py:49, in as_tensor_variable(x, name, ndim, **kwargs)
     17 def as_tensor_variable(
     18     x: TensorLike, name: Optional[str] = None, ndim: Optional[int] = None, **kwargs
     19 ) -> "TensorVariable":
     20     """Convert `x` into an equivalent `TensorVariable`.
     21 
     22     This function can be used to turn ndarrays, numbers, `ScalarType` instances,
   (...)
     47 
     48     """
---> 49     return _as_tensor_variable(x, name, ndim, **kwargs)

File ~/anaconda3/envs/custom_39_pymc/lib/python3.9/functools.py:888, in singledispatch.<locals>.wrapper(*args, **kw)
    884 if not args:
    885     raise TypeError(f'{funcname} requires at least '
    886                     '1 positional argument')
--> 888 return dispatch(args[0].__class__)(*args, **kw)

File ~/SageMaker/custom_39_venv/lib/python3.9/site-packages/aesara/tensor/__init__.py:56, in _as_tensor_variable(x, name, ndim, **kwargs)
     52 @singledispatch
     53 def _as_tensor_variable(
     54     x: TensorLike, name: Optional[str], ndim: Optional[int], **kwargs
     55 ) -> "TensorVariable":
---> 56     raise NotImplementedError(f"Cannot convert {x!r} to a tensor variable.")

NotImplementedError: Cannot convert Elemwise{pow,no_inplace}.0 to a tensor variable.

Previously (in-line with what worked in my PyMC3 model), I used

w = at.as_tensor_variable([at.power(alpha,i) for i in range(L)])

instead of

w = at.power(alpha, np.arange(L))

But that also spat out errors that I’m unable to resolve.

Any suggestions or advice are greatly appreciated.

PyMC version: 5.0.2
Aesara version: 2.8.10
NumPy version: 1.24.2

Running on Python 3.9.16 on a AWS SageMaker Notebook instance

With PyMC 5 you have to use PyTensor and not Aesara.

Just replace import aesara.tensor by import pytensor.tensor.

1 Like