Skip to content

Provide a function / recipe to cache endpoint outputs #455

@dionhaefner

Description

@dionhaefner

Summary

Examples:

  • Cache forward passes to use during backwards pass.
  • Compute value_and_grad during apply and return cached values when calling AD endpoints.
  • Cache full Jacobians and use them for JVP / VJPs.
  • More??

Why is this needed?

Sometimes it is cheaper to re-use computations. People building Tesseracts should have the freedom to exploit this.

Usage example

E.g. provide an example that uses functools.lru_cache or similar, and that provides a hash function for pytrees (potentially as part of tesseract_core.runtime).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions