Dynamic RV shape or choosing from it

Questions version agnostic
2

I have written this code:

def main():
    data = generate_data()

    import pymc as pm
    import pytensor.tensor as at
    import arviz as az

    indexes = np.indices((users_count, len(dialog_prob)))[1]

    with pm.Model(coords={'user': range(users_count), 'num': range(len(dialog_prob))}) as m:
        dialog_duration = pm.Exponential('', lam=2, dims=('user', 'num'))
        dialog_count = pm.Categorical('dialog_count', p=dialog_prob, shape=users_count, dims='user')

        dialog_duration = pm.Deterministic('all_duration',
                                           at.where(indexes < dialog_count[:, None], dialog_duration, 0),
                                           dims=('user', 'num'))
        user_duration = pm.Deterministic('user_duration', at.sum(dialog_duration, axis=1), dims='user')

        trace = pm.sample_prior_predictive(random_seed=42)

    print(az.summary(trace))
    breakpoint()
    print()

This code samples correct sum by users, but i don’t know how to reshape dialog_duration to matrix like this:

[
[user_0, dialog_0_len],
[user_0, dialog_1_len],
[user_1, dialog_0_len],
[user_4, dialog_0_len],
]

is it possible? Or maybe exists some another way to sample data like this?