address review comments

Author: jpbrodrick89

examples/vectoradd_jax/tesseract_api.py

@@ -10,7 +10,7 @@
 from typing_extensions import Self
 
 from tesseract_core.runtime import Array, Differentiable, Float32
-from tesseract_core.runtime.gradient_recipes import (
+from tesseract_core.runtime.gradient_endpoints.jax_recipes import (
     jax_abstract_eval,
     jax_apply,
     jax_jacobian,

tesseract_core/runtime/__init__.py

@@ -14,11 +14,12 @@
             continue
         if not (path.parent / "__init__.py").exists():
             continue
-        # gradient_recipes/ is an opt-in plug-in surface — each recipe brings
-        # its own framework dep (e.g. jax/equinox). Tesseracts that opt into
-        # a recipe install its deps via tesseract_requirements.txt; tesseracts
-        # that don't shouldn't be required to.
-        if "gradient_recipes" in path.parts:
+        # gradient_endpoints/ is an opt-in plug-in surface — each backend
+        # (jax_recipes, future torch_recipes, ...) brings its own framework
+        # dep. Tesseracts that opt into a backend install its deps via
+        # tesseract_requirements.txt; tesseracts that don't shouldn't be
+        # required to.
+        if "gradient_endpoints" in path.parts:
             continue
 
         package_path = ".".join(

…act_core/runtime/gradient_recipes/jax.py …adient_endpoints/jax_recipes/__init__.pytesseract_core/runtime/gradient_recipes/jax.py renamed to tesseract_core/runtime/gradient_endpoints/jax_recipes/__init__.py tesseract_core/runtime/gradient_recipes/jax.py renamed to tesseract_core/runtime/gradient_endpoints/jax_recipes/__init__.py

@@ -6,20 +6,21 @@
 These functions remove the boilerplate from a JAX-backed Tesseract
 ``tesseract_api.py`` by providing one-line implementations of the
 ``apply``, ``jacobian``, ``jacobian_vector_product``,
-``vector_jacobian_product`` and ``abstract_eval`` endpoints, along with
-shared VJP residual caching across ``apply`` and ``vector_jacobian_product``.
-
-Cache behaviour
----------------
-By default :func:`jax_apply` runs ``jax.vjp`` internally and stashes the
-resulting backward function in :data:`jax_vjp_cache`. A subsequent
-:func:`jax_vjp` call on the same inputs reuses that backward, skipping the
-redundant forward pass — typically a 10-20% saving on moderate-to-deep
-networks. Call :func:`set_jax_cache_size` with ``0`` to disable caching
-entirely; both :func:`jax_apply` and :func:`jax_vjp` then bypass the cache
-machinery and run their forward passes directly.
+``vector_jacobian_product`` and ``abstract_eval`` endpoints. The user
+supplies a single ``apply_jit`` (already wrapped in ``@eqx.filter_jit``);
+the helpers compose the JAX transforms around it.
+
+VJP residual caching is enabled by default: :func:`jax_apply` stashes the
+backward function it constructs and :func:`jax_vjp` reuses it on the next
+matching call — typically a ~10-20% speedup on the standard tesseract-jax
+``apply → vjp`` pattern (which is hit by ``jax.value_and_grad``,
+``jax.jacrev``, and even plain ``jax.grad`` since tesseract-jax's
+``custom_vjp.fwd`` always runs ``apply()`` first). Call
+:func:`set_jax_vjp_cache_size` with ``0`` to disable. See its docstring
+for the full taxonomy.
 """
 
+import hashlib
 from collections.abc import Callable
 from typing import Any
 
@@ -32,42 +33,90 @@
     LRUCache,
     filter_func,
     flatten_with_paths,
-    hash_pytree_leaves,
     set_at_path,
 )
 
-# Default size of the VJP residual cache. Use :func:`set_jax_cache_size` to
-# change at runtime; do not mutate this directly.
-_jax_cache_size = 1
-
 #: Module-level VJP residual cache. ``jax_apply`` populates it; ``jax_vjp``
-#: reads from it.
-jax_vjp_cache = LRUCache(maxsize=_jax_cache_size)
-
+#: reads from it. Set to ``None`` (via :func:`set_jax_vjp_cache_size` with
+#: ``0``) to disable caching — both helpers then bypass the cache machinery
+#: entirely.
+jax_vjp_cache: LRUCache | None = LRUCache(maxsize=1)
+
+
+def set_jax_vjp_cache_size(size: int) -> None:
+    """Enable or resize the VJP residual cache, discarding existing entries.
+
+    When :func:`jax_apply` and :func:`jax_vjp` are called in sequence on the
+    same inputs (the typical pattern under tesseract-jax's ``custom_vjp``,
+    including under ``jax.value_and_grad`` and ``jax.jacrev``), caching the
+    backward function produced by :func:`jax_apply`'s internal ``jax.vjp``
+    call lets the subsequent :func:`jax_vjp` skip the redundant forward
+    pass. Typical speedup ~10-20% on moderate-to-deep networks.
+
+    Best fit:
+      - tesseract-jax workflows that go through ``custom_vjp`` (every
+        gradient evaluation runs ``apply()`` first via ``fwd``, so the
+        cache reliably hits on the subsequent ``vjp()`` — true even
+        under plain ``jax.grad``).
+      - Manual ``apply()`` followed by multiple ``vjp()`` calls on the
+        same inputs (e.g. iterative solvers, CG-style inverse problems).
+
+    Poor fit:
+      - Very small models where the forward pass is microseconds: the
+        cache machinery overhead exceeds the saved work.
+
+    Args:
+        size: Number of cache slots. ``1`` (the default) covers the standard
+            apply → vjp pattern. ``0`` disables caching entirely (both
+            :func:`jax_apply` and :func:`jax_vjp` then bypass the cache
+            machinery). Increase for workflows that interleave multiple
+            ``apply()`` calls before their corresponding ``vjp()`` calls.
+    """
+    global jax_vjp_cache
+    jax_vjp_cache = LRUCache(maxsize=size) if size > 0 else None
 
-def set_jax_cache_size(size: int) -> None:
-    """Resize the VJP cache. Discards any existing entries."""
-    global _jax_cache_size, jax_vjp_cache
-    _jax_cache_size = size
-    jax_vjp_cache = LRUCache(maxsize=size)
 
+def hash_tree(tree: Any) -> bytes:
+    """Compute a SHA-256 digest over a pytree's structure and leaves.
 
-def _hash_inputs(inputs_dict: dict) -> bytes:
-    """Compute a SHA-256 hash of a pytree of inputs for cache key comparison."""
-    leaves, treedef = jax.tree.flatten(inputs_dict)
-    return hash_pytree_leaves(leaves, treedef)
+    Suitable for use as an :class:`LRUCache` key. For array leaves, the digest
+    incorporates dtype + shape + raw bytes so leaves with identical bytes but
+    different interpretations (e.g. ``int64[4]`` vs ``int64[2,2]``) don't
+    collide. Scalar and primitive leaves use Python's ``hash()`` except for
+    ``str``/``bytes`` (whose hash is randomized across processes by Python's
+    PYTHONHASHSEED — we encode them directly instead).
+    """
+    leaves, treedef = jax.tree.flatten(tree)
+    h = hashlib.sha256()
+    h.update(str(treedef).encode())
+    for leaf in leaves:
+        if hasattr(leaf, "tobytes"):
+            h.update(str(leaf.dtype).encode())
+            h.update(str(leaf.shape).encode())
+            h.update(leaf.tobytes())
+        elif isinstance(leaf, (str, bytes)):
+            h.update(leaf if isinstance(leaf, bytes) else leaf.encode())
+        else:
+            h.update(hash(leaf).to_bytes(8, "big", signed=True))
+    return h.digest()
 
 
 def jax_apply(apply_jit: Callable, inputs: BaseModel) -> dict:
-    """Run ``apply_jit`` and populate the VJP cache.
+    """Run ``apply_jit`` and, if caching is enabled, populate the VJP cache.
+
+    ``apply_jit`` is assumed to already be JIT-compiled (e.g. wrapped with
+    ``@eqx.filter_jit``); this helper does not jit it. The user-facing
+    ``apply`` endpoint may want to do pre/post-processing around the call,
+    so we cannot wrap it in a jit internally.
 
-    A subsequent :func:`jax_vjp` call on the same input reuses the cached
-    backward. When caching is disabled (see :func:`set_jax_cache_size`),
-    this is just ``apply_jit(inputs.model_dump())``.
+    When :data:`jax_vjp_cache` is set (see :func:`set_jax_vjp_cache_size`),
+    the forward pass is run via ``jax.vjp`` so the resulting backward
+    function can be stashed and reused by a later :func:`jax_vjp` call.
+    Otherwise this is just ``apply_jit(inputs.model_dump())``.
     """
     inputs_dict = inputs.model_dump()
 
-    if _jax_cache_size <= 0:
+    if jax_vjp_cache is None:
         return apply_jit(inputs_dict)
 
     # Compute forward pass via jax.vjp to cache residuals for a potential
@@ -85,7 +134,7 @@ def _apply_for_vjp(inputs_dict: dict) -> tuple:
     out = eqx.combine(diff_primals, static_primals)
 
     cotangent_template = jax.tree.map(jnp.zeros_like, diff_primals)
-    jax_vjp_cache.put(_hash_inputs(inputs_dict), (vjp_func, cotangent_template))
+    jax_vjp_cache.put(hash_tree(inputs_dict), (vjp_func, cotangent_template))
     return out
 
 
@@ -98,8 +147,11 @@ def jax_vjp(
 ) -> dict[str, Any]:
     """Compute the vector-Jacobian product.
 
-    Reuses the cached backward from a prior :func:`jax_apply` call if one is
-    available; otherwise falls through to a fresh ``jax.vjp`` evaluation.
+    Reuses the cached backward from a prior :func:`jax_apply` call when one
+    is available (see :func:`set_jax_vjp_cache_size`); otherwise falls
+    through to a freshly JIT-compiled ``jax.vjp`` evaluation. The JIT
+    compilation happens internally on the first miss for a given
+    (input shape/dtype, path subset) combination and is cached for reuse.
     """
     inputs_dict = inputs.model_dump()
 
@@ -108,8 +160,8 @@ def jax_vjp(
     # value_and_grad is followed by jax.jacrev, which decomposes into many
     # vjp calls per output basis vector.
     if (
-        _jax_cache_size > 0
-        and (cached := jax_vjp_cache.get(_hash_inputs(inputs_dict))) is not None
+        jax_vjp_cache is not None
+        and (cached := jax_vjp_cache.get(hash_tree(inputs_dict))) is not None
     ):
         vjp_func, cotangent_template = cached
         full_cotangent = jax.tree.map(jnp.zeros_like, cotangent_template)
@@ -134,7 +186,11 @@ def jax_jvp(
     jvp_outputs: set[str],
     tangent_vector: dict[str, Any],
 ) -> dict[str, Any]:
-    """Compute the Jacobian-vector product via :func:`jax.jvp`."""
+    """Compute the Jacobian-vector product via :func:`jax.jvp`.
+
+    JIT compilation is applied internally and cached per
+    ``(input shape/dtype, jvp_inputs, jvp_outputs)`` combination.
+    """
     return _jvp_jit(
         apply_jit,
         inputs.model_dump(),
@@ -150,7 +206,11 @@ def jax_jacobian(
     jac_inputs: set[str],
     jac_outputs: set[str],
 ) -> dict[str, dict[str, Any]]:
-    """Compute the Jacobian via :func:`jax.jacrev`."""
+    """Compute the Jacobian via :func:`jax.jacrev`.
+
+    JIT compilation is applied internally and cached per
+    ``(input shape/dtype, jac_inputs, jac_outputs)`` combination.
+    """
     return _jac_jit(
         apply_jit, inputs.model_dump(), tuple(jac_inputs), tuple(jac_outputs)
     )

tesseract_core/runtime/gradient_recipes/__init__.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import collections
-import hashlib
 import re
 import threading
 from collections.abc import Callable, Iterable, Mapping, Sequence
@@ -199,8 +198,7 @@ def put(self, key: bytes, value: Any) -> None:
             return
         with self._lock:
             if key in self._cache:
-                if len(self._cache) > 1 and next(reversed(self._cache)) != key:
-                    self._cache.move_to_end(key)
+                self._cache.move_to_end(key)
             self._cache[key] = value
             while len(self._cache) > self._maxsize:
                 self._cache.popitem(last=False)
@@ -210,37 +208,11 @@ def get(self, key: bytes) -> Any | None:
         with self._lock:
             if key not in self._cache:
                 return None
-            if len(self._cache) > 1 and next(reversed(self._cache)) != key:
-                self._cache.move_to_end(key)
+            self._cache.move_to_end(key)
             return self._cache[key]
 
-    def pop(self, key: bytes) -> Any | None:
-        """Remove and return the value for *key*, or ``None`` on a miss."""
-        with self._lock:
-            return self._cache.pop(key, None)
-
     @property
     def size(self) -> int:
         """Return the number of entries currently in the cache."""
         with self._lock:
             return len(self._cache)
-
-
-def hash_pytree_leaves(leaves: Iterable, treedef: Any) -> bytes:
-    """Compute a SHA-256 digest over the leaves and structure of a pytree.
-
-    Args:
-        leaves: Flat sequence of leaf values (arrays or scalars).
-        treedef: Tree definition object (its ``str()`` is hashed to capture structure).
-
-    Returns:
-        A 32-byte SHA-256 digest suitable for use as an :class:`LRUCache` key.
-    """
-    h = hashlib.sha256()
-    h.update(str(treedef).encode())
-    for leaf in leaves:
-        if hasattr(leaf, "tobytes"):
-            h.update(leaf.tobytes())
-        else:
-            h.update(str(leaf).encode())
-    return h.digest()

tesseract_core/runtime/tree_transforms.py

@@ -10,45 +10,23 @@
 from pydantic import BaseModel
 
 from tesseract_core.runtime import Differentiable, Float32
-from tesseract_core.runtime.gradient_recipes import (
+from tesseract_core.runtime.gradient_endpoints.jax_recipes import (
     jax_abstract_eval,
     jax_apply,
     jax_jacobian,
     jax_jvp,
     jax_vjp,
 )
 
+# VJP residual caching is enabled by default and gives a ~10-20% speedup on
+# the typical tesseract-jax apply → vjp pattern. See the docstring of
+# `set_jax_vjp_cache_size` for the full when-it-helps taxonomy. To disable
+# or resize, uncomment below:
 #
-# Cache configuration
-#
-# apply() runs jax.vjp internally and caches the resulting backward function.
-# A subsequent vector_jacobian_product call on the same inputs reuses that
-# backward, skipping the redundant forward pass that vjp would otherwise repeat.
-#
-# Under typical tesseract-jax usage (gradient-based code calling the Tesseract
-# through tesseract-jax's custom_vjp), apply() runs first on every gradient
-# evaluation — even under plain jax.grad — so the cache reliably hits on the
-# subsequent vjp call. This makes the recipe close to a free win for any
-# gradient-based workflow: Adam/SGD, L-BFGS / line search, value_and_grad-style
-# code, etc. Typical speedup is ~10-20% on moderate-to-deep networks.
-#
-# The cache also helps when a single apply is followed by multiple vjp calls
-# on the same input — for example, jax.jacrev (which decomposes into one vjp
-# per output basis vector) or CG-style inverse-problem solvers (each inner CG
-# step computes J^T u at the same iterate). Each subsequent vjp on the same
-# input reuses the cached residuals.
-#
-# Poor fit:
-#   - Very small models where the forward pass is microseconds: the cache
-#     machinery overhead exceeds the saved work.
-#
-# Uncomment to disable caching entirely (apply() and vjp() then bypass the
-# cache machinery and run their forward passes directly), or to increase the
-# size for workloads that interleave several apply() calls before their
-# corresponding vjp() calls.
-#
-# from tesseract_core.runtime.gradient_recipes import set_jax_cache_size
-# set_jax_cache_size(0)
+# from tesseract_core.runtime.gradient_endpoints.jax_recipes import (
+#     set_jax_vjp_cache_size,
+# )
+# set_jax_vjp_cache_size(0)  # disable
 
 #
 # Schemata