Skip to content

Allow jacobian and hessian of arbitrary dimension expressions #1225

Open
@ricardoV94

Description

@ricardoV94

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?

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions