pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Vectorize node should return list of variables not node

Open ricardoV94 opened this issue 1 year ago • 0 comments

Description

vectorize_node implicitly assumes that whenever we want to vectorize a node, we will return a new node that has a 1-to-1 mapping with the original outputs, but this is too restrictive. It could be the case we want to vectorize a single node with two variables coming from different nodes, or a single output from a multi-valued node. There's no reason why we need a one node -> one node mapping.

https://github.com/pymc-devs/pytensor/blob/79ff97a54bfdd98aa877cc6a04b5a4e3782cfbd6/pytensor/graph/replace.py#L208-L211

For backwards compatibility we should check if the returned object is an Apply and issue a warning that this form is deprecated (but still use it) and instead a list of outputs (like in the rewrites) should be returned. All our implementations in PyTensor should switch to returning a list of variables.

The catch/warning could be done here: https://github.com/pymc-devs/pytensor/blob/79ff97a54bfdd98aa877cc6a04b5a4e3782cfbd6/pytensor/graph/replace.py#L214-L217

Then everything that calls vectorize_node should now expect a list as output. Like here: https://github.com/pymc-devs/pytensor/blob/79ff97a54bfdd98aa877cc6a04b5a4e3782cfbd6/pytensor/graph/replace.py#L301-L308

ricardoV94 avatar Jul 08 '24 11:07 ricardoV94