PyTensor, but it would be good if you could simplify it quite a bit and make it a standalone reproducible snippet