Skip to content

Slow performance of jax.jacobian #170

@jpbrodrick89

Description

@jpbrodrick89

Even if available Jax is unable to call a Tesseract's jacobian endpoint and always chooses to batch over jvp (jacfwd) or (vjp). As Tesseracts don't currently support batched tangents and unbatched primals (see issue in tesseract-core for more details on design options) Tesseract-JAX is forced to perform multiple sequential calls to the Tesseract which can be very slow and inefficient. Until we support batched tangent vectors in Tesseract-core there is nothing we can do here but noting this here for awareness.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions