pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Allow jacobian and hessian of arbitrary dimension expressions

Open ricardoV94 opened this issue 10 months ago • 0 comments

Description

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

x = pt.vector("x", shape=(3,))
y = pt.outer(x[1:], x[2:])
assert y.type.shape == (2, 1)
try:
    jacobian(y, x)
except Exception as exc:
    print(exc)  # jacobian expects a 1 dimensional variable as `expression`. If not use flatten to make it a vector

jac_y = jacobian(y.ravel(), x).reshape((*y.shape, *x.shape))
assert jac_y.type.shape == (2, 1, 3)

I don't see why we can't do the ravel -> reshape for the users? JAX accepts non-vector jacobian just fine.

The hessian is trickier as it requires also the combinations of the inputs?

ricardoV94 avatar Feb 19 '25 18:02 ricardoV94