Implement vectorized jacobian and allow arbitrary expression shapes

[ricardoV94]

Seeing about a 3.5x, (2x for larger x) time speedup for this trivial case:

from pytensor.gradient import jacobian
from pytensor import function
import pytensor.tensor as pt

x = pt.vector("x", shape=(3,))
y = pt.outer(x, x)
print(y.type.shape)

jac_y = jacobian(y, x, vectorize=False)
print(jac_y.type.shape)

fn = function([x], jac_y, profile=True)
fn.dprint(print_type=True)

%timeit fn([0, 1, 2])
fn([0, 1, 2])

Memory footprint will grow though, specially if intermediate operations are much larger than the final jacobian. Also not all graphs will be safely vectorizable, so I would leave the Scan option as a default for a while.

Related Issue

• Related to Add option to implement vectorized hessian/jacobian #756
• Related to Allow jacobian and hessian of arbitrary dimension expressions #1225

[ricardoV94]

This is equivalent, but we sidestep the implicit eye that shows up in the gradient of expr[idx] and pass it to the Lop directly. Otherwise there is a Blockwise(IncSubtensor) that is not pretty