pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Implement vectorized jacobian and allow arbitrary expression shapes

Open ricardoV94 opened this issue 10 months ago • 0 comments

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

ricardoV94 avatar Feb 20 '25 09:02 ricardoV94