What is the best way to vectorize an operation on tensor in PyMC?

Questions
1

Hello everyone, the title says it all,

For more context : I built a numpy model to compute fake images, which relies on the utilization of either the np.meshgrid function or a @np.vectorize decorator to enable summation over a higher dimension, which is instantly reduced. I struggle to find any equivalent of these processes for tensor variable, and there is no equivalent of meshgrid in aesara.tensor (as mgrid and ogrid are not suited for tensor variable).

Do you have any idea?

2

Can you share or explain what exactly your numpy function is doing?

3

An example of what I want to do with vectorize

import numpy as np 

@np.vectorize
def func(x):
    y = np.linspace(-1, 1, 10)
    return np.sum((x**2 + y**2)**(1/2))

x = np.linspace(-1, 1, 10)
res = func(x)

or with meshgrid

x = np.linspace(-1, 1, 10)
y = np.linspace(-1, 1, 10)

X, Y = np.meshgrid(x,y,indexing='ij')

res = np.sum( (X**2 + Y**2)**(1/2), axis=-1)

However, in my PyMC model, I would like to perform similar operations with x and y being 1d-tensor variables directly dependent on RV parameters

4

The best solution I found yet relies on broadcasting

X = x[:,None]
Y = y[None,:]
res = (X**2 + Y**2)**(1/2)

This piece of code worked with x and y being aesara tensors unlike the previous meshgrid and vectorize solution.

1 Like