pytensor
pytensor copied to clipboard
Implement vectorized jacobian and allow arbitrary expression shapes
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
- [x] Related to #756
- [x] Related to #1225