wrapper for more complex switch statement

[drbenvincent]

So I watched Austin Rochford's into talk where he showed a simple example of using Aesara with the Monte Hall problem...

import pymc as pm
from aesara import tensor as aet
import numpy as np

with pm.Model() as monty_model:
    prize = pm.DiscreteUniform("prize", 0, 2, initval=0)
    
    p_open = aet.switch(
        aet.eq(prize, 0), 
        np.array([0, 0.5, 0.5]),  # prize behind 1st door
        aet.switch(
            aet.eq(prize, 1),
            np.array([0, 0, 1]),  # prize behind 2nd door
            np.array([0, 1, 0])  # prize behind 3rd door
        )
    )
    
    pm.Categorical("opened", p_open, observed=2)
    
    trace = pm.sample(10_000)

for i in [0, 1, 2]:
    print(f'P(prize behind door {i}|opened=2) = {(trace.posterior.prize==i).mean().data}')

# P(prize behind door 0|opened=2) = 0.326225
# P(prize behind door 1|opened=2) = 0.673775
# P(prize behind door 2|opened=2) = 0.0

So I'm wondering if there could be a more more user-friendly function that could take the place of that nested aet.switch? So you could have an API like this

p_open = aet.multi_switch(prize, 
                          np.array([0, 0.5, 0.5]),
                          np.array([0, 0, 1]),
                          np.array([0, 1, 0]))

or

p_open = aet.multi_switch(prize, 
                          [np.array([0, 0.5, 0.5]),
                           np.array([0, 0, 1]),
                           np.array([0, 1, 0])])

So aet.multi_switch (pick better name) would just be syntactic sugar for the nested switch.