`ifelse` does not work with arbitrary data structures

[rlouf]

What I expect

I expect ifelse to work when I pass two arbitrarily nested python data structures that have identical structures, and whose leaves are TensorVariables of identical shape and dtype.

This would greatly facilitate the use of aesara for "complex" projects such as aehmc.

What I observe

It works with a tuple of TensorVariable

An IfElse Op can be created passing a tuple or a list of TensorVariables for each branch of the condition:

import aesara.tensor as at
from aesara.ifelse import ifelse

cond = at.as_tensor(0, dtype=bool)

q_left = at.vector()
energy_left = at.scalar()
energy_grad_left = at.vector()
state_left = (q_left, energy_left, energy_grad_left)

q_right = at.vector()
energy_right = at.scalar()
energy_grad_right = at.vector()
state_right = (q_right, energy_right, energy_grad_right)

new_state = ifelse(cond, state_left, state_right)

It does not work with nested tuples

However, it does not work with nested structures. Here is a simplified version of something I have encountered in aehmc:

weight_left = at.scalar()
left = (state_left, weight_left)

weight_right = at.scalar()
right = (state_right, weight_right)

new = ifelse(cond, left, right)

Which returns a TypeError.

It does not work with dictionaries

In the same way, it will not work with dictionaries as input, the following also returns a TypeError:

import aesara.tensor as at
from aesara.tensor.var import TensorVariable
from aesara.ifelse import ifelse

state_left = {"q": at.vector(), "energy": at.scalar(), "energy_grad": at.vector()}
state_right = {"q": at.vector(), "energy": at.scalar(), "energy_grad": at.vector()}

cond = at.as_tensor(0, dtype=bool)
state = ifelse(cond, state_left, state_right)

It does not work with namedtuples

IfElse returns a tuple when NamedTuples are passed an input, causing errors downstream.

from typing import NamedTuple

import aesara.tensor as at
from aesara.tensor.var import TensorVariable
from aesara.ifelse import ifelse


class State(NamedTuple):
    q: TensorVariable
    energy: TensorVariable
    energy_grad: TensorVariable

state_left = State(at.vector(), at.scalar(), at.vector())
state_right = State(at.vector(), at.scalar(), at.vector())

cond = at.as_tensor(0, dtype=bool)
state = ifelse(cond, state_left, state_right)
state.q

[brandonwillard]

Most of what you've mentioned has little to do with ifelse, and has primarily to do with tuple support. In other words, generic tuples aren't supported (i.e. modeled at the graph-level), and that's really what this issue is requesting.

Aesara's type system doesn't have the same flexibility as Python (e.g. there is no Tuple[int, float, str] type). You'll find similar issues in Numba and other transpilation libraries for related reasons.

As we've discussed before, there is a named tuple-like data structure that could provide some of these conveniences. As a matter of fact, the kind of conveniences you request could possibly be covered by custom Types and/or Variables; ones with implementations that are restricted to their needs. The challenges largely come from making such things work for everything, everywhere, and everyhow.

Remember, there's still a legacy C backend that's used by default, so that needs to be appeased—i.e. if you want to use those "tuples" at that level, you'll need some generic struct-generating code.

If one only wants to make a user or even dev-level interface simpler, they should use their own tuple-like containers and concatenate then * them into function calls. Likewise, a developer can make their own simplfied interfaces that accept these tuple-like objects, or custom classes for that matter, and construct graphs with those.

For example, see my comments here and here about such an improved interface in AeHMC.

In the end, there is a mild convenience to having Ops that can return and take tuples, but all that's really providing is a sometimes convenient nesting structure that is just as well handled by simple indexing—regarding only the things that pertain to the content of this issue, of course.

Also, I don't see the outputs of your examples, so it's not clear exactly how/why all these things aren't working.