Test suite failing in latest version of aeppl

[ricardoV94]

I have opened an issue in aesara-devs/aeppl#92

It is still unclear what is going on, but seems related to transforms

[brandonwillard]

pymc/tests/test_dist_math.py::test_multinomial_bound might be failing due to a change in a Shape sub-graph. Printing aesara.dprint(modelA.logpt), the change between AePPL versions appears to take the node

 |   |     | |     | |Shape [id DB] ''
 |   |     | |     | | |Elemwise{sub,no_inplace} [id DC] ''
 |   |     | |     | |   |Subtensor{:int64:} [id DD] ''
 |   |     | |     | |   | |Elemwise{log,no_inplace} [id DE] ''
 |   |     | |     | |   | | |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   |     | |     | |   | |ScalarConstant{-1} [id DF]
 |   |     | |     | |   |Elemwise{true_div,no_inplace} [id DG] ''
 |   |     | |     | |     |InplaceDimShuffle{x} [id DH] ''
 |   |     | |     | |     | |Sum{axis=[0], acc_dtype=float64} [id DI] ''
 |   |     | |     | |     |   |Elemwise{log,no_inplace} [id DE] ''
 |   |     | |     | |     |InplaceDimShuffle{x} [id DJ] ''
 |   |     | |     | |       |Subtensor{int64} [id DK] ''
 |   |     | |     | |         |Shape [id DL] ''
 |   |     | |     | |         | |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   |     | |     | |         |ScalarConstant{-1} [id DM]

and changes it to

 |   |     | |     | |Shape [id DC] ''
 |   |     | |     | | |p_simplex__ [id T]

Here are the full graph print-outs for reference:

old graph

Elemwise{add,no_inplace} [id A] '__logp'
 |Sum{acc_dtype=float64} [id B] ''
 | |Elemwise{mul,no_inplace} [id C] ''
 |   |Elemwise{add,no_inplace} [id D] 'p_simplex___logprob'
 |   | |Elemwise{switch,no_inplace} [id E] ''
 |   | | |All [id F] ''
 |   | | | |MakeVector{dtype='bool'} [id G] ''
 |   | | |   |All [id H] ''
 |   | | |   | |Elemwise{neq,no_inplace} [id I] ''
 |   | | |   |   |Elemwise{mul,no_inplace} [id J] ''
 |   | | |   |   | |TensorConstant{1} [id K]
 |   | | |   |   | |All [id L] ''
 |   | | |   |   |   |Elemwise{ge,no_inplace} [id M] ''
 |   | | |   |   |     |TransformedVariable [id N] 'p_simplex___simplex'
 |   | | |   |   |     | |Elemwise{true_div,no_inplace} [id O] ''
 |   | | |   |   |     | | |Elemwise{exp,no_inplace} [id P] ''
 |   | | |   |   |     | | | |Elemwise{sub,no_inplace} [id Q] ''
 |   | | |   |   |     | | |   |Join [id R] ''
 |   | | |   |   |     | | |   | |TensorConstant{0} [id S]
 |   | | |   |   |     | | |   | |p_simplex__ [id T]
 |   | | |   |   |     | | |   | |Elemwise{neg,no_inplace} [id U] ''
 |   | | |   |   |     | | |   |   |InplaceDimShuffle{x} [id V] ''
 |   | | |   |   |     | | |   |     |Sum{axis=[0], acc_dtype=float64} [id W] ''
 |   | | |   |   |     | | |   |       |p_simplex__ [id T]
 |   | | |   |   |     | | |   |InplaceDimShuffle{x} [id X] ''
 |   | | |   |   |     | | |     |MaxAndArgmax{axis=(0,)}.0 [id Y] 'max'
 |   | | |   |   |     | | |       |Join [id R] ''
 |   | | |   |   |     | | |InplaceDimShuffle{x} [id Z] ''
 |   | | |   |   |     | |   |Sum{axis=[0], acc_dtype=float64} [id BA] ''
 |   | | |   |   |     | |     |Elemwise{exp,no_inplace} [id P] ''
 |   | | |   |   |     | |p_simplex__ [id T]
 |   | | |   |   |     |InplaceDimShuffle{x} [id BB] ''
 |   | | |   |   |       |TensorConstant{0} [id BC]
 |   | | |   |   |TensorConstant{0} [id BD]
 |   | | |   |All [id BE] ''
 |   | | |   | |Elemwise{neq,no_inplace} [id BF] ''
 |   | | |   |   |Elemwise{mul,no_inplace} [id BG] ''
 |   | | |   |   | |TensorConstant{1} [id BH]
 |   | | |   |   | |All [id BI] ''
 |   | | |   |   |   |Elemwise{le,no_inplace} [id BJ] ''
 |   | | |   |   |     |TransformedVariable [id N] 'p_simplex___simplex'
 |   | | |   |   |     |InplaceDimShuffle{x} [id BK] ''
 |   | | |   |   |       |TensorConstant{1} [id BL]
 |   | | |   |   |TensorConstant{0} [id BM]
 |   | | |   |All [id BN] ''
 |   | | |     |Elemwise{neq,no_inplace} [id BO] ''
 |   | | |       |Elemwise{mul,no_inplace} [id BP] ''
 |   | | |       | |TensorConstant{1} [id BQ]
 |   | | |       | |All [id BR] ''
 |   | | |       |   |Elemwise{gt,no_inplace} [id BS] ''
 |   | | |       |     |TensorConstant{(2,) of 1.0} [id BT]
 |   | | |       |     |InplaceDimShuffle{x} [id BU] ''
 |   | | |       |       |TensorConstant{0} [id BV]
 |   | | |       |TensorConstant{0} [id BW]
 |   | | |Elemwise{add,no_inplace} [id BX] ''
 |   | | | |Sum{axis=[0], acc_dtype=float64} [id BY] ''
 |   | | | | |Elemwise{sub,no_inplace} [id BZ] ''
 |   | | | |   |Elemwise{switch,no_inplace} [id CA] ''
 |   | | | |   | |Elemwise{eq,no_inplace} [id CB] ''
 |   | | | |   | | |TransformedVariable [id N] 'p_simplex___simplex'
 |   | | | |   | | |InplaceDimShuffle{x} [id CC] ''
 |   | | | |   | |   |TensorConstant{0} [id CD]
 |   | | | |   | |Elemwise{switch,no_inplace} [id CE] ''
 |   | | | |   | | |Elemwise{eq,no_inplace} [id CF] ''
 |   | | | |   | | | |Elemwise{sub,no_inplace} [id CG] ''
 |   | | | |   | | | | |TensorConstant{(2,) of 1.0} [id BT]
 |   | | | |   | | | | |InplaceDimShuffle{x} [id CH] ''
 |   | | | |   | | | |   |TensorConstant{1} [id CI]
 |   | | | |   | | | |InplaceDimShuffle{x} [id CJ] ''
 |   | | | |   | | |   |TensorConstant{0} [id CK]
 |   | | | |   | | |InplaceDimShuffle{x} [id CL] ''
 |   | | | |   | | | |TensorConstant{0.0} [id CM]
 |   | | | |   | | |InplaceDimShuffle{x} [id CN] ''
 |   | | | |   | |   |TensorConstant{-inf} [id CO]
 |   | | | |   | |Elemwise{mul,no_inplace} [id CP] ''
 |   | | | |   |   |Elemwise{sub,no_inplace} [id CG] ''
 |   | | | |   |   |Elemwise{log,no_inplace} [id CQ] ''
 |   | | | |   |     |TransformedVariable [id N] 'p_simplex___simplex'
 |   | | | |   |Elemwise{gammaln,no_inplace} [id CR] ''
 |   | | | |     |TensorConstant{(2,) of 1.0} [id BT]
 |   | | | |Elemwise{gammaln,no_inplace} [id CS] ''
 |   | | |   |Sum{axis=[0], acc_dtype=float64} [id CT] ''
 |   | | |     |TensorConstant{(2,) of 1.0} [id BT]
 |   | | |TensorConstant{-inf} [id CU]
 |   | |Sum{axis=[0], acc_dtype=float64} [id CV] ''
 |   |   |Elemwise{sub,no_inplace} [id CW] ''
 |   |     |Elemwise{add,no_inplace} [id CX] ''
 |   |     | |InplaceDimShuffle{x} [id CY] ''
 |   |     | | |Elemwise{log,no_inplace} [id CZ] ''
 |   |     | |   |Elemwise{add,no_inplace} [id DA] ''
 |   |     | |     |Subtensor{int64} [id DB] ''
 |   |     | |     | |Shape [id DC] ''
 |   |     | |     | | |p_simplex__ [id T]
 |   |     | |     | |ScalarConstant{-1} [id DD]
 |   |     | |     |TensorConstant{1} [id DE]
 |   |     | |Elemwise{mul,no_inplace} [id DF] ''
 |   |     |   |InplaceDimShuffle{x} [id DG] ''
 |   |     |   | |Elemwise{add,no_inplace} [id DA] ''
 |   |     |   |InplaceDimShuffle{x} [id DH] ''
 |   |     |     |Sum{axis=[0], acc_dtype=float64} [id DI] ''
 |   |     |       |p_simplex__ [id T]
 |   |     |Elemwise{mul,no_inplace} [id DJ] ''
 |   |       |InplaceDimShuffle{x} [id DK] ''
 |   |       | |Elemwise{add,no_inplace} [id DA] ''
 |   |       |Elemwise{log,no_inplace} [id DL] ''
 |   |         |InplaceDimShuffle{x} [id DM] ''
 |   |           |Sum{axis=[0], acc_dtype=float64} [id DN] ''
 |   |             |Elemwise{exp,no_inplace} [id DO] ''
 |   |               |Join [id DP] ''
 |   |                 |TensorConstant{0} [id DQ]
 |   |                 |Elemwise{add,no_inplace} [id DR] ''
 |   |                 | |p_simplex__ [id T]
 |   |                 | |InplaceDimShuffle{x} [id DH] ''
 |   |                 |Alloc [id DS] ''
 |   |                   |TensorConstant{0.0} [id DT]
 |   |                   |Subtensor{int64} [id DU] ''
 |   |                     |Shape [id DV] ''
 |   |                     | |InplaceDimShuffle{x} [id DH] ''
 |   |                     |ScalarConstant{0} [id DW]
 |   |TensorConstant{1.0} [id DX]
 |Elemwise{add,no_inplace} [id DY] ''
   |Sum{acc_dtype=float64} [id DZ] ''
   | |Elemwise{mul,no_inplace} [id EA] ''
   |   |Elemwise{switch,no_inplace} [id EB] 'x_logprob'
   |   | |All [id EC] ''
   |   | | |MakeVector{dtype='bool'} [id ED] ''
   |   | |   |All [id EE] ''
   |   | |   | |Elemwise{neq,no_inplace} [id EF] ''
   |   | |   |   |Elemwise{mul,no_inplace} [id EG] ''
   |   | |   |   | |TensorConstant{1} [id EH]
   |   | |   |   | |All [id EI] ''
   |   | |   |   |   |Elemwise{ge,no_inplace} [id EJ] ''
   |   | |   |   |     |TensorConstant{[1 5]} [id EK]
   |   | |   |   |     |InplaceDimShuffle{x} [id EL] ''
   |   | |   |   |       |TensorConstant{0} [id EM]
   |   | |   |   |TensorConstant{0} [id EN]
   |   | |   |All [id EO] ''
   |   | |   | |Elemwise{neq,no_inplace} [id EP] ''
   |   | |   |   |Elemwise{mul,no_inplace} [id EQ] ''
   |   | |   |   | |TensorConstant{1} [id ER]
   |   | |   |   | |All [id ES] ''
   |   | |   |   |   |Elemwise{ge,no_inplace} [id ET] ''
   |   | |   |   |     |p_simplex__ [id T]
   |   | |   |   |     |InplaceDimShuffle{x} [id EU] ''
   |   | |   |   |       |TensorConstant{0} [id EV]
   |   | |   |   |TensorConstant{0} [id EW]
   |   | |   |All [id EX] ''
   |   | |   | |Elemwise{neq,no_inplace} [id EY] ''
   |   | |   |   |Elemwise{mul,no_inplace} [id EZ] ''
   |   | |   |   | |TensorConstant{1} [id FA]
   |   | |   |   | |All [id FB] ''
   |   | |   |   |   |Elemwise{le,no_inplace} [id FC] ''
   |   | |   |   |     |p_simplex__ [id T]
   |   | |   |   |     |InplaceDimShuffle{x} [id FD] ''
   |   | |   |   |       |TensorConstant{1} [id FE]
   |   | |   |   |TensorConstant{0} [id FF]
   |   | |   |All [id FG] ''
   |   | |     |Elemwise{neq,no_inplace} [id FH] ''
   |   | |       |Elemwise{mul,no_inplace} [id FI] ''
   |   | |       | |TensorConstant{1} [id FJ]
   |   | |       | |Elemwise{or_,no_inplace} [id FK] ''
   |   | |       |   |Elemwise{and_,no_inplace} [id FL] ''
   |   | |       |   | |Elemwise{le,no_inplace} [id FM] ''
   |   | |       |   | | |Elemwise{abs,no_inplace} [id FN] ''
   |   | |       |   | | | |Elemwise{sub,no_inplace} [id FO] ''
   |   | |       |   | | |   |Sum{acc_dtype=float64} [id FP] ''
   |   | |       |   | | |   | |p_simplex__ [id T]
   |   | |       |   | | |   |TensorConstant{1} [id FQ]
   |   | |       |   | | |TensorConstant{1.001e-05} [id FR]
   |   | |       |   | |Elemwise{invert,no_inplace} [id FS] ''
   |   | |       |   |   |Elemwise{or_,no_inplace} [id FT] ''
   |   | |       |   |     |Elemwise{or_,no_inplace} [id FU] ''
   |   | |       |   |     | |Elemwise{isnan,no_inplace} [id FV] ''
   |   | |       |   |     | | |Sum{acc_dtype=float64} [id FP] ''
   |   | |       |   |     | |Alloc [id FW] ''
   |   | |       |   |     |   |TensorConstant{False} [id FX]
   |   | |       |   |     |Elemwise{or_,no_inplace} [id FY] ''
   |   | |       |   |       |Elemwise{isinf,no_inplace} [id FZ] ''
   |   | |       |   |       | |Sum{acc_dtype=float64} [id FP] ''
   |   | |       |   |       |Alloc [id GA] ''
   |   | |       |   |         |TensorConstant{False} [id GB]
   |   | |       |   |Elemwise{and_,no_inplace} [id GC] ''
   |   | |       |     |Elemwise{and_,no_inplace} [id GD] ''
   |   | |       |     | |Elemwise{isinf,no_inplace} [id FZ] ''
   |   | |       |     | |Alloc [id GA] ''
   |   | |       |     |Elemwise{eq,no_inplace} [id GE] ''
   |   | |       |       |Elemwise{mul,no_inplace} [id GF] ''
   |   | |       |       | |Elemwise{isinf,no_inplace} [id FZ] ''
   |   | |       |       | |Elemwise{sgn,no_inplace} [id GG] ''
   |   | |       |       |   |Sum{acc_dtype=float64} [id FP] ''
   |   | |       |       |Elemwise{mul,no_inplace} [id GH] ''
   |   | |       |         |Alloc [id GA] ''
   |   | |       |         |Elemwise{sgn,no_inplace} [id GI] ''
   |   | |       |           |TensorConstant{1} [id GJ]
   |   | |       |TensorConstant{0} [id GK]
   |   | |Elemwise{add,no_inplace} [id GL] ''
   |   | | |Elemwise{sub,no_inplace} [id GM] ''
   |   | | | |Elemwise{gammaln,no_inplace} [id GN] ''
   |   | | | | |Elemwise{add,no_inplace} [id GO] ''
   |   | | | |   |TensorConstant{6} [id GP]
   |   | | | |   |TensorConstant{1} [id GQ]
   |   | | | |Sum{acc_dtype=float64} [id GR] ''
   |   | | |   |Elemwise{gammaln,no_inplace} [id GS] ''
   |   | | |     |Elemwise{add,no_inplace} [id GT] ''
   |   | | |       |TensorConstant{[1 5]} [id EK]
   |   | | |       |InplaceDimShuffle{x} [id GU] ''
   |   | | |         |TensorConstant{1} [id GV]
   |   | | |Sum{acc_dtype=float64} [id GW] ''
   |   | |   |Elemwise{mul,no_inplace} [id GX] ''
   |   | |     |TensorConstant{[1 5]} [id EK]
   |   | |     |Elemwise{log,no_inplace} [id GY] ''
   |   | |       |p_simplex__ [id T]
   |   | |TensorConstant{-inf} [id GZ]
   |   |TensorConstant{1.0} [id HA]
   |TensorConstant{0} [id HB]

new graph

Elemwise{add,no_inplace} [id A] '__logp'
 |Sum{acc_dtype=float64} [id B] ''
 | |Elemwise{mul,no_inplace} [id C] ''
 |   |Elemwise{add,no_inplace} [id D] 'p_simplex___logprob'
 |   | |Elemwise{switch,no_inplace} [id E] ''
 |   | | |All [id F] ''
 |   | | | |MakeVector{dtype='bool'} [id G] ''
 |   | | |   |All [id H] ''
 |   | | |   | |Elemwise{neq,no_inplace} [id I] ''
 |   | | |   |   |Elemwise{mul,no_inplace} [id J] ''
 |   | | |   |   | |TensorConstant{1} [id K]
 |   | | |   |   | |All [id L] ''
 |   | | |   |   |   |Elemwise{ge,no_inplace} [id M] ''
 |   | | |   |   |     |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   | | |   |   |     | |Elemwise{exp,no_inplace} [id O] ''
 |   | | |   |   |     | | |Elemwise{sub,no_inplace} [id P] ''
 |   | | |   |   |     | |   |Join [id Q] ''
 |   | | |   |   |     | |   | |TensorConstant{0} [id R]
 |   | | |   |   |     | |   | |p_simplex__ [id S]
 |   | | |   |   |     | |   | |Elemwise{neg,no_inplace} [id T] ''
 |   | | |   |   |     | |   |   |InplaceDimShuffle{x} [id U] ''
 |   | | |   |   |     | |   |     |Sum{axis=[0], acc_dtype=float64} [id V] ''
 |   | | |   |   |     | |   |       |p_simplex__ [id S]
 |   | | |   |   |     | |   |InplaceDimShuffle{x} [id W] ''
 |   | | |   |   |     | |     |MaxAndArgmax{axis=(0,)}.0 [id X] 'max'
 |   | | |   |   |     | |       |Join [id Q] ''
 |   | | |   |   |     | |InplaceDimShuffle{x} [id Y] ''
 |   | | |   |   |     |   |Sum{axis=[0], acc_dtype=float64} [id Z] ''
 |   | | |   |   |     |     |Elemwise{exp,no_inplace} [id O] ''
 |   | | |   |   |     |InplaceDimShuffle{x} [id BA] ''
 |   | | |   |   |       |TensorConstant{0} [id BB]
 |   | | |   |   |TensorConstant{0} [id BC]
 |   | | |   |All [id BD] ''
 |   | | |   | |Elemwise{neq,no_inplace} [id BE] ''
 |   | | |   |   |Elemwise{mul,no_inplace} [id BF] ''
 |   | | |   |   | |TensorConstant{1} [id BG]
 |   | | |   |   | |All [id BH] ''
 |   | | |   |   |   |Elemwise{le,no_inplace} [id BI] ''
 |   | | |   |   |     |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   | | |   |   |     |InplaceDimShuffle{x} [id BJ] ''
 |   | | |   |   |       |TensorConstant{1} [id BK]
 |   | | |   |   |TensorConstant{0} [id BL]
 |   | | |   |All [id BM] ''
 |   | | |     |Elemwise{neq,no_inplace} [id BN] ''
 |   | | |       |Elemwise{mul,no_inplace} [id BO] ''
 |   | | |       | |TensorConstant{1} [id BP]
 |   | | |       | |All [id BQ] ''
 |   | | |       |   |Elemwise{gt,no_inplace} [id BR] ''
 |   | | |       |     |TensorConstant{(2,) of 1.0} [id BS]
 |   | | |       |     |InplaceDimShuffle{x} [id BT] ''
 |   | | |       |       |TensorConstant{0} [id BU]
 |   | | |       |TensorConstant{0} [id BV]
 |   | | |Elemwise{add,no_inplace} [id BW] ''
 |   | | | |Sum{axis=[0], acc_dtype=float64} [id BX] ''
 |   | | | | |Elemwise{sub,no_inplace} [id BY] ''
 |   | | | |   |Elemwise{switch,no_inplace} [id BZ] ''
 |   | | | |   | |Elemwise{eq,no_inplace} [id CA] ''
 |   | | | |   | | |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   | | | |   | | |InplaceDimShuffle{x} [id CB] ''
 |   | | | |   | |   |TensorConstant{0} [id CC]
 |   | | | |   | |Elemwise{switch,no_inplace} [id CD] ''
 |   | | | |   | | |Elemwise{eq,no_inplace} [id CE] ''
 |   | | | |   | | | |Elemwise{sub,no_inplace} [id CF] ''
 |   | | | |   | | | | |TensorConstant{(2,) of 1.0} [id BS]
 |   | | | |   | | | | |InplaceDimShuffle{x} [id CG] ''
 |   | | | |   | | | |   |TensorConstant{1} [id CH]
 |   | | | |   | | | |InplaceDimShuffle{x} [id CI] ''
 |   | | | |   | | |   |TensorConstant{0} [id CJ]
 |   | | | |   | | |InplaceDimShuffle{x} [id CK] ''
 |   | | | |   | | | |TensorConstant{0.0} [id CL]
 |   | | | |   | | |InplaceDimShuffle{x} [id CM] ''
 |   | | | |   | |   |TensorConstant{-inf} [id CN]
 |   | | | |   | |Elemwise{mul,no_inplace} [id CO] ''
 |   | | | |   |   |Elemwise{sub,no_inplace} [id CF] ''
 |   | | | |   |   |Elemwise{log,no_inplace} [id CP] ''
 |   | | | |   |     |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   | | | |   |Elemwise{gammaln,no_inplace} [id CQ] ''
 |   | | | |     |TensorConstant{(2,) of 1.0} [id BS]
 |   | | | |Elemwise{gammaln,no_inplace} [id CR] ''
 |   | | |   |Sum{axis=[0], acc_dtype=float64} [id CS] ''
 |   | | |     |TensorConstant{(2,) of 1.0} [id BS]
 |   | | |TensorConstant{-inf} [id CT]
 |   | |Sum{axis=[0], acc_dtype=float64} [id CU] ''
 |   |   |Elemwise{sub,no_inplace} [id CV] ''
 |   |     |Elemwise{add,no_inplace} [id CW] ''
 |   |     | |InplaceDimShuffle{x} [id CX] ''
 |   |     | | |Elemwise{log,no_inplace} [id CY] ''
 |   |     | |   |Elemwise{add,no_inplace} [id CZ] ''
 |   |     | |     |Subtensor{int64} [id DA] ''
 |   |     | |     | |Shape [id DB] ''
 |   |     | |     | | |Elemwise{sub,no_inplace} [id DC] ''
 |   |     | |     | |   |Subtensor{:int64:} [id DD] ''
 |   |     | |     | |   | |Elemwise{log,no_inplace} [id DE] ''
 |   |     | |     | |   | | |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   |     | |     | |   | |ScalarConstant{-1} [id DF]
 |   |     | |     | |   |Elemwise{true_div,no_inplace} [id DG] ''
 |   |     | |     | |     |InplaceDimShuffle{x} [id DH] ''
 |   |     | |     | |     | |Sum{axis=[0], acc_dtype=float64} [id DI] ''
 |   |     | |     | |     |   |Elemwise{log,no_inplace} [id DE] ''
 |   |     | |     | |     |InplaceDimShuffle{x} [id DJ] ''
 |   |     | |     | |       |Subtensor{int64} [id DK] ''
 |   |     | |     | |         |Shape [id DL] ''
 |   |     | |     | |         | |Elemwise{true_div,no_inplace} [id N] 'p_simplex___simplex'
 |   |     | |     | |         |ScalarConstant{-1} [id DM]
 |   |     | |     | |ScalarConstant{-1} [id DN]
 |   |     | |     |TensorConstant{1} [id DO]
 |   |     | |Elemwise{mul,no_inplace} [id DP] ''
 |   |     |   |InplaceDimShuffle{x} [id DQ] ''
 |   |     |   | |Elemwise{add,no_inplace} [id CZ] ''
 |   |     |   |InplaceDimShuffle{x} [id DR] ''
 |   |     |     |Sum{axis=[0], acc_dtype=float64} [id DS] ''
 |   |     |       |Elemwise{sub,no_inplace} [id DC] ''
 |   |     |Elemwise{mul,no_inplace} [id DT] ''
 |   |       |InplaceDimShuffle{x} [id DU] ''
 |   |       | |Elemwise{add,no_inplace} [id CZ] ''
 |   |       |Elemwise{log,no_inplace} [id DV] ''
 |   |         |InplaceDimShuffle{x} [id DW] ''
 |   |           |Sum{axis=[0], acc_dtype=float64} [id DX] ''
 |   |             |Elemwise{exp,no_inplace} [id DY] ''
 |   |               |Join [id DZ] ''
 |   |                 |TensorConstant{0} [id EA]
 |   |                 |Elemwise{add,no_inplace} [id EB] ''
 |   |                 | |Elemwise{sub,no_inplace} [id DC] ''
 |   |                 | |InplaceDimShuffle{x} [id DR] ''
 |   |                 |Alloc [id EC] ''
 |   |                   |TensorConstant{0.0} [id ED]
 |   |                   |Subtensor{int64} [id EE] ''
 |   |                     |Shape [id EF] ''
 |   |                     | |InplaceDimShuffle{x} [id DR] ''
 |   |                     |ScalarConstant{0} [id EG]
 |   |TensorConstant{1.0} [id EH]
 |Elemwise{add,no_inplace} [id EI] ''
   |Sum{acc_dtype=float64} [id EJ] ''
   | |Elemwise{mul,no_inplace} [id EK] ''
   |   |Elemwise{switch,no_inplace} [id EL] 'x_logprob'
   |   | |All [id EM] ''
   |   | | |MakeVector{dtype='bool'} [id EN] ''
   |   | |   |All [id EO] ''
   |   | |   | |Elemwise{neq,no_inplace} [id EP] ''
   |   | |   |   |Elemwise{mul,no_inplace} [id EQ] ''
   |   | |   |   | |TensorConstant{1} [id ER]
   |   | |   |   | |All [id ES] ''
   |   | |   |   |   |Elemwise{ge,no_inplace} [id ET] ''
   |   | |   |   |     |TensorConstant{[1 5]} [id EU]
   |   | |   |   |     |InplaceDimShuffle{x} [id EV] ''
   |   | |   |   |       |TensorConstant{0} [id EW]
   |   | |   |   |TensorConstant{0} [id EX]
   |   | |   |All [id EY] ''
   |   | |   | |Elemwise{neq,no_inplace} [id EZ] ''
   |   | |   |   |Elemwise{mul,no_inplace} [id FA] ''
   |   | |   |   | |TensorConstant{1} [id FB]
   |   | |   |   | |All [id FC] ''
   |   | |   |   |   |Elemwise{ge,no_inplace} [id FD] ''
   |   | |   |   |     |Elemwise{true_div,no_inplace} [id FE] 'p_simplex___simplex'
   |   | |   |   |     | |Elemwise{exp,no_inplace} [id FF] ''
   |   | |   |   |     | | |Elemwise{sub,no_inplace} [id FG] ''
   |   | |   |   |     | |   |Join [id FH] ''
   |   | |   |   |     | |   | |TensorConstant{0} [id FI]
   |   | |   |   |     | |   | |p_simplex__ [id S]
   |   | |   |   |     | |   | |Elemwise{neg,no_inplace} [id FJ] ''
   |   | |   |   |     | |   |   |InplaceDimShuffle{x} [id FK] ''
   |   | |   |   |     | |   |     |Sum{axis=[0], acc_dtype=float64} [id FL] ''
   |   | |   |   |     | |   |       |p_simplex__ [id S]
   |   | |   |   |     | |   |InplaceDimShuffle{x} [id FM] ''
   |   | |   |   |     | |     |MaxAndArgmax{axis=(0,)}.0 [id FN] 'max'
   |   | |   |   |     | |       |Join [id FH] ''
   |   | |   |   |     | |InplaceDimShuffle{x} [id FO] ''
   |   | |   |   |     |   |Sum{axis=[0], acc_dtype=float64} [id FP] ''
   |   | |   |   |     |     |Elemwise{exp,no_inplace} [id FF] ''
   |   | |   |   |     |InplaceDimShuffle{x} [id FQ] ''
   |   | |   |   |       |TensorConstant{0} [id FR]
   |   | |   |   |TensorConstant{0} [id FS]
   |   | |   |All [id FT] ''
   |   | |   | |Elemwise{neq,no_inplace} [id FU] ''
   |   | |   |   |Elemwise{mul,no_inplace} [id FV] ''
   |   | |   |   | |TensorConstant{1} [id FW]
   |   | |   |   | |All [id FX] ''
   |   | |   |   |   |Elemwise{le,no_inplace} [id FY] ''
   |   | |   |   |     |Elemwise{true_div,no_inplace} [id FE] 'p_simplex___simplex'
   |   | |   |   |     |InplaceDimShuffle{x} [id FZ] ''
   |   | |   |   |       |TensorConstant{1} [id GA]
   |   | |   |   |TensorConstant{0} [id GB]
   |   | |   |All [id GC] ''
   |   | |     |Elemwise{neq,no_inplace} [id GD] ''
   |   | |       |Elemwise{mul,no_inplace} [id GE] ''
   |   | |       | |TensorConstant{1} [id GF]
   |   | |       | |Elemwise{or_,no_inplace} [id GG] ''
   |   | |       |   |Elemwise{and_,no_inplace} [id GH] ''
   |   | |       |   | |Elemwise{le,no_inplace} [id GI] ''
   |   | |       |   | | |Elemwise{abs,no_inplace} [id GJ] ''
   |   | |       |   | | | |Elemwise{sub,no_inplace} [id GK] ''
   |   | |       |   | | |   |Sum{acc_dtype=float64} [id GL] ''
   |   | |       |   | | |   | |Elemwise{true_div,no_inplace} [id FE] 'p_simplex___simplex'
   |   | |       |   | | |   |TensorConstant{1} [id GM]
   |   | |       |   | | |TensorConstant{1.001e-05} [id GN]
   |   | |       |   | |Elemwise{invert,no_inplace} [id GO] ''
   |   | |       |   |   |Elemwise{or_,no_inplace} [id GP] ''
   |   | |       |   |     |Elemwise{or_,no_inplace} [id GQ] ''
   |   | |       |   |     | |Elemwise{isnan,no_inplace} [id GR] ''
   |   | |       |   |     | | |Sum{acc_dtype=float64} [id GL] ''
   |   | |       |   |     | |Alloc [id GS] ''
   |   | |       |   |     |   |TensorConstant{False} [id GT]
   |   | |       |   |     |Elemwise{or_,no_inplace} [id GU] ''
   |   | |       |   |       |Elemwise{isinf,no_inplace} [id GV] ''
   |   | |       |   |       | |Sum{acc_dtype=float64} [id GL] ''
   |   | |       |   |       |Alloc [id GW] ''
   |   | |       |   |         |TensorConstant{False} [id GX]
   |   | |       |   |Elemwise{and_,no_inplace} [id GY] ''
   |   | |       |     |Elemwise{and_,no_inplace} [id GZ] ''
   |   | |       |     | |Elemwise{isinf,no_inplace} [id GV] ''
   |   | |       |     | |Alloc [id GW] ''
   |   | |       |     |Elemwise{eq,no_inplace} [id HA] ''
   |   | |       |       |Elemwise{mul,no_inplace} [id HB] ''
   |   | |       |       | |Elemwise{isinf,no_inplace} [id GV] ''
   |   | |       |       | |Elemwise{sgn,no_inplace} [id HC] ''
   |   | |       |       |   |Sum{acc_dtype=float64} [id GL] ''
   |   | |       |       |Elemwise{mul,no_inplace} [id HD] ''
   |   | |       |         |Alloc [id GW] ''
   |   | |       |         |Elemwise{sgn,no_inplace} [id HE] ''
   |   | |       |           |TensorConstant{1} [id HF]
   |   | |       |TensorConstant{0} [id HG]
   |   | |Elemwise{add,no_inplace} [id HH] ''
   |   | | |Elemwise{sub,no_inplace} [id HI] ''
   |   | | | |Elemwise{gammaln,no_inplace} [id HJ] ''
   |   | | | | |Elemwise{add,no_inplace} [id HK] ''
   |   | | | |   |TensorConstant{6} [id HL]
   |   | | | |   |TensorConstant{1} [id HM]
   |   | | | |Sum{acc_dtype=float64} [id HN] ''
   |   | | |   |Elemwise{gammaln,no_inplace} [id HO] ''
   |   | | |     |Elemwise{add,no_inplace} [id HP] ''
   |   | | |       |TensorConstant{[1 5]} [id EU]
   |   | | |       |InplaceDimShuffle{x} [id HQ] ''
   |   | | |         |TensorConstant{1} [id HR]
   |   | | |Sum{acc_dtype=float64} [id HS] ''
   |   | |   |Elemwise{mul,no_inplace} [id HT] ''
   |   | |     |TensorConstant{[1 5]} [id EU]
   |   | |     |Elemwise{log,no_inplace} [id HU] ''
   |   | |       |Elemwise{true_div,no_inplace} [id FE] 'p_simplex___simplex'
   |   | |TensorConstant{-inf} [id HV]
   |   |TensorConstant{1.0} [id HW]
   |TensorConstant{0} [id HX]

[ricardoV94]

Yeah it is a transform replacement issue, which in this case also affects the shape of the variables: aesara-devs/aeppl#92 (comment)