pymc-devs/pytensor

Description I found this project: https://github.com/guillaume-osmo/mlx-addons It looks vibe coded, but it has MLX GPU kernels for cholesky, QR, random SVD, LU, triangular solve, and others. Also far be it from me to judge vibe coding these days :) Would be a big speedup for many models, especially GP/KF/MvN stuff. It's MIT licensed so we could "take inspiration" or try to work with them directly. Tagging the dev @guillaume-osmo/@thegodone in case you're interested in collaborating on this.

We've talked about this before so I just want an issue to track it. The idea is that: perform should be in python, but it should be as simple as possible mode = 'py' should have python-specific dispatches where we can put mode full-featured "performance" python code. An example would be the QZ Op. All that should be moved out of perform, perform should just do scipy.linalg.qz Replace the default FAST_COMPILE with python mode (no compile is really fast right?) There might not actually be much to do to get this off the ground, just some refactor to make things nice and hide complexity in certain cases. Not sure about how this would impact random variables, scan, etc.

Description Integer-to-integer casts on shape arguments of Alloc (and other shape-consuming ops like Reshape) are not removed during canonicalization, even though these ops accept any integer dtype for shape inputs. For example, Alloc(val, int32_var.astype("int64")) keeps the Cast{int64} node in the optimized graph, despite Alloc working correctly with int32 shape inputs directly. Reproducer import pytensor.tensor as pt from pytensor import function, dprint from pytensor.compile.mode import get_default_mode x = pt.scalar('x') s = pt.scalar('s', dtype='int32') g1 = pt.alloc(x, s) g2 = pt.alloc(x, s.astype('int64')) mode = get_default_mode().including('canonicalize') f1 = function([x, s], g1, mode=mode) f2 = function([x, s], g2, mode=mode) print('Without cast:') dprint(f1) # Alloc [id A] # |-- x [id B] # `-- s [id C] print('With useless int32->int64 cast:') dprint(f2) # Alloc [id A] # |-- x [id B] # `-- Cast{int64} [id C] # `-- s [id D] The same issue applies to all integer

Description https://github.com/pymc-devs/pytensor/actions/runs/22731848114/job/65922904767?pr=1937#step:6:222 Started failing with jax 0.9.1

Description After #1756 Advanced and IncSubtensor take idx_list (just like Subtensor did). It would be nice to use the same string representation we used for Subtensor to make dprint more readable. import pytensor.tensor as pt x = pt.zeros((2, 3, 4, 5)) basic_sub = x[0, :, :1, 1] basic_sub.dprint() # Subtensor{i, :, :stop, j} [id A] # ├─ Alloc [id B] # │ ├─ 0.0 [id C] # │ ├─ 2 [id D] # │ ├─ 3 [id E] # │ ├─ 4 [id F] # │ └─ 5 [id G] # ├─ 0 [id H] # ├─ 1 [id I] # └─ 1 [id J] adv_sub = x[[0], :, :1, 1] adv_sub.dprint() # AdvancedSubtensor{idx_list=(0, slice(None, None, None), slice(None, 1, None), 2)} [id A] # ├─ Alloc [id B] # │ ├─ 0.0 [id C] # │ ├─ 2 [id D] # │ ├─ 3 [id E] # │ ├─ 4 [id F] # │ └─ 5 [id G] # ├─ [0] [id H] # ├─ TensorFromScalar [id I] # │ └─ 1 [id J] # └─ TensorFromScalar [id K] # └─ 1 [id L] (We should also check what's going on with th