Version dependant slowing down of Gaussian Mixture sampling in Ubuntu 20.04

I have also seen the following:

5.9.0
 │     │     │     │  │           │  ├─ Cholesky{lower=True, destructive=False, on_error='nan'} [id EE] 73

5.9.1
│     │     │     │  │        │  ├─ ExpandDims{axis=0} [id EJ] 91
│     │     │     │  │        │  │  └─ Cholesky{lower=True, destructive=False, on_error='nan'} [id EF] 73

but mainly most of the differences are of the form

5.9.0
 │  │  │  │  │     │  │     │  │     │     └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2} [id KR] 295

5.9.1
 │  │  │  │  │     │  │     │  │     │     └─ Blockwise{SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=1}, (m,m),(m)->(m)} [id KS] 294

One another big difference seems to be here

5.9.0

│  │  │  │  │  │  │  │  │        └─ Composite{((i0 * i1) - i2)} [id MA] 580
│  │  │  │  │  │  │  │  │           ├─ Transpose{axes=[1, 0]} [id MB] 388
│  │  │  │  │  │  │  │  │           │  └─ dot [id MC] 380
│  │  │  │  │  │  │  │  │           │     ├─ Switch [id LS] 176
│  │  │  │  │  │  │  │  │           │     │  └─ ···
│  │  │  │  │  │  │  │  │           │     └─ Composite{switch(i4, (i3 + switch(i2, (-1.0 * i0 * i1), 0.0)), 1)} [id MD] 373
│  │  │  │  │  │  │  │  │           │        ├─ dot [id ME] 298
│  │  │  │  │  │  │  │  │           │        │  ├─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2} [id JL] 287
│  │  │  │  │  │  │  │  │           │        │  │  └─ ···
│  │  │  │  │  │  │  │  │           │        │  └─ Transpose{axes=[1, 0]} [id CC] 205
│  │  │  │  │  │  │  │  │           │        │     └─ ···
│  │  │  │  │  │  │  │  │           │        ├─ Tri{dtype='float64'} [id MF] 366
│  │  │  │  │  │  │  │  │           │        │  ├─ Subtensor{i} [id MG] 354
│  │  │  │  │  │  │  │  │           │        │  │  ├─ Subtensor{start:} [id MH] 342
│  │  │  │  │  │  │  │  │           │        │  │  │  ├─ Shape [id MI] 330
│  │  │  │  │  │  │  │  │           │        │  │  │  │  └─ Neg [id MJ] 314
│  │  │  │  │  │  │  │  │           │        │  │  │  │     └─ dot [id ME] 298
│  │  │  │  │  │  │  │  │           │        │  │  │  │        └─ ···
│  │  │  │  │  │  │  │  │           │        │  │  │  └─ ScalarFromTensor [id MK] 12
│  │  │  │  │  │  │  │  │           │        │  │  │     └─ -2 [id ML]
│  │  │  │  │  │  │  │  │           │        │  │  └─ ScalarFromTensor [id MM] 13
│  │  │  │  │  │  │  │  │           │        │  │     └─ 0 [id MN]
│  │  │  │  │  │  │  │  │           │        │  ├─ Subtensor{i} [id MO] 353
│  │  │  │  │  │  │  │  │           │        │  │  ├─ Subtensor{start:} [id MH] 342
│  │  │  │  │  │  │  │  │           │        │  │  │  └─ ···
│  │  │  │  │  │  │  │  │           │        │  │  └─ ScalarFromTensor [id MP] 11
│  │  │  │  │  │  │  │  │           │        │  │     └─ 1 [id MQ]
│  │  │  │  │  │  │  │  │           │        │  └─ 0 [id E]

5.9.1
 └─ Composite{((i0 * i1) - i2)} [id MB] 604
 │  │  │  │  │  │  │  │  │           ├─ Transpose{axes=[1, 0]} [id MC] 413
 │  │  │  │  │  │  │  │  │           │  └─ dot [id MD] 404
 │  │  │  │  │  │  │  │  │           │     ├─ Switch [id LS] 167
 │  │  │  │  │  │  │  │  │           │     │  └─ ···
 │  │  │  │  │  │  │  │  │           │     └─ Composite{switch(i3, (i2 + switch(i1, (-i0), 0.0)), 1)} [id ME] 397
 │  │  │  │  │  │  │  │  │           │        ├─ Sum{axis=0} [id MF] 391
 │  │  │  │  │  │  │  │  │           │        │  └─ Mul [id MG] 384
 │  │  │  │  │  │  │  │  │           │        │     ├─ Blockwise{dot, (i00,i01),(i10,i11)->(o00,o01)} [id MH] 313
 │  │  │  │  │  │  │  │  │           │        │     │  ├─ ExpandDims{axis=2} [id MI] 297
 │  │  │  │  │  │  │  │  │           │        │     │  │  └─ Blockwise{SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=1}, (m,m),(m)->(m)} [id JQ] 286
 │  │  │  │  │  │  │  │  │           │        │     │  │     └─ ···
 │  │  │  │  │  │  │  │  │           │        │     │  └─ ExpandDims{axis=1} [id MJ] 220
 │  │  │  │  │  │  │  │  │           │        │     │     └─ Blockwise{SolveTriangular{trans=0, unit_diagonal=False, lower=True, check_finite=True, b_ndim=1}, (m,m),(m)->(m)} [id CC] 207
 │  │  │  │  │  │  │  │  │           │        │     │        └─ ···
 │  │  │  │  │  │  │  │  │           │        │     └─ Blockwise{Tri{dtype='float64'}, (),(),()->(o00,o01)} [id MK] 377
 │  │  │  │  │  │  │  │  │           │        │        ├─ Blockwise{Subtensor{i}, (i00),()->()} [id ML] 365
 │  │  │  │  │  │  │  │  │           │        │        │  ├─ Blockwise{Subtensor{start:}, (i00),()->(o00)} [id MM] 353
 │  │  │  │  │  │  │  │  │           │        │        │  │  ├─ Blockwise{Shape, (i00,i01)->(o00)} [id MN] 341
 │  │  │  │  │  │  │  │  │           │        │        │  │  │  └─ Neg [id MO] 329
 │  │  │  │  │  │  │  │  │           │        │        │  │  │     └─ Blockwise{dot, (i00,i01),(i10,i11)->(o00,o01)} [id MH] 313
 │  │  │  │  │  │  │  │  │           │        │        │  │  │        └─ ···
 │  │  │  │  │  │  │  │  │           │        │        │  │  └─ [-2] [id MP]
 │  │  │  │  │  │  │  │  │           │        │        │  └─ [0] [id MQ]
 │  │  │  │  │  │  │  │  │           │        │        ├─ Blockwise{Subtensor{i}, (i00),()->()} [id MR] 364
 │  │  │  │  │  │  │  │  │           │        │        │  ├─ Blockwise{Subtensor{start:}, (i00),()->(o00)} [id MM] 353
 │  │  │  │  │  │  │  │  │           │        │        │  │  └─ ···
 │  │  │  │  │  │  │  │  │           │        │        │  └─ [1] [id MS]
 │  │  │  │  │  │  │  │  │           │        │        └─ [0] [id CO]

this is also different

5.9.0
 │  │  │  │  │  │  │  └─ Subtensor{i} [id YM] 822
 │  │  │  │  │  │  │     ├─ Shape [id YN] 796
 │  │  │  │  │  │  │     │  └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id YO] 764
 │  │  │  │  │  │  │     │     └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2} [id YP] 725

5.9.1
├─ ARange{dtype='int64'} [id ZY] 865
│  │  │  │  │  │  │  ├─ 0 [id E]
│  │  │  │  │  │  │  ├─ Subtensor{i} [id ZZ] 839
│  │  │  │  │  │  │  │  ├─ Shape [id BAA] 813
│  │  │  │  │  │  │  │  │  └─ ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id BAB] 781
│  │  │  │  │  │  │  │  │     └─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, 

There is also one part where Dot was replaced with Sum and Mult (I guess in the end they do the same thing but still…) but then it is followed by something called Blockwise dot which seemed strange. And also the inputs look some what different in switch (i4,i3,i2,i1 vs i3,i2,i1).

5.9.0
└─ Composite{switch(i4, (i3 + switch(i2, (-1.0 * i0 * i1), 0.0)), 1)} [id BAN] 377
 │  │  │  │  │           │        ├─ dot [id BAO] 307
 │  │  │  │  │           │        │  ├─ SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=2} [id KK] 293
 │  │  │  │  │           │        │  │  └─ ···
 │  │  │  │  │           │        │  └─ Transpose{axes=[1, 0]} [id EQ] 202
 │  │  │  │  │           │        │     └─ ···

5.9.1
└─ Composite{switch(i3, (i2 + switch(i1, (-i0), 0.0)), 1)} [id BBD] 402
 │  │  │  │  │           │        ├─ Sum{axis=0} [id BBE] 395
 │  │  │  │  │           │        │  └─ Mul [id BBF] 388
 │  │  │  │  │           │        │     ├─ Blockwise{dot, (i00,i01),(i10,i11)->(o00,o01)} [id BBG] 322
 │  │  │  │  │           │        │     │  ├─ ExpandDims{axis=2} [id BBH] 306
 │  │  │  │  │           │        │     │  │  └─ Blockwise{SolveTriangular{trans=0, unit_diagonal=False, lower=False, check_finite=True, b_ndim=1}, (m,m),(m)->(m)} [id KM] 292

Possibly another related difference (note All{axis=0} vs All{axis=None})

5.9.0
│     │     │     │  │  └─ CAReduce{Composite{(i0 + sqr(i1))}, axis=1} [id CB] 213
│     │     │     │  │     └─ Transpose{axes=[1, 0]} [id CC] 205
│     │     │     │  │        └─ SolveTriangular{trans=0, unit_diagonal=False, lower=True, check_finite=True, b_ndim=2} [id CD] 197
│     │     │     │  │           ├─ Switch [id CE] 156
│     │     │     │  │           │  ├─ ExpandDims{axes=[0, 1]} [id CF] 138
│     │     │     │  │           │  │  └─ All{axes=None} [id CG] 120
│     │     │     │  │           │  │     └─ Gt [id CH] 101

5.9.1
│     │     │     │  │  └─ CAReduce{Composite{(i0 + sqr(i1))}, axis=1} [id CB] 219
│     │     │     │  │     └─ Blockwise{SolveTriangular{trans=0, unit_diagonal=False, lower=True, check_finite=True, b_ndim=1}, (m,m),(m)->(m)} [id CC] 207
│     │     │     │  │        ├─ Switch [id CD] 166
│     │     │     │  │        │  ├─ ExpandDims{axes=[0, 1, 2]} [id CE] 145
│     │     │     │  │        │  │  └─ All{axis=0} [id CF] 125
│     │     │     │  │        │  │     └─ Gt [id CG] 106