reproduce error

Author: andrinr

tests/conftest.py

@@ -20,6 +20,8 @@ def get_tesseract_folders():
     tesseract_folders = [
         "univariate_tesseract",
         "nested_tesseract",
+        "non_abstract_tesseract",
+        "vectoradd_tesseract",
         # Add more as needed
     ]
     return tesseract_folders
@@ -86,3 +88,4 @@ def served_tesseract():
 served_univariate_tesseract_raw = make_tesseract_fixture("univariate_tesseract")
 served_nested_tesseract_raw = make_tesseract_fixture("nested_tesseract")
 served_non_abstract_tesseract = make_tesseract_fixture("non_abstract_tesseract")
+served_vectoradd_tesseract = make_tesseract_fixture("vectoradd_tesseract")

tests/test_endtoend.py

@@ -3,6 +3,7 @@
 
 
 import jax
+import jax.numpy as jnp
 import numpy as np
 import pytest
 from jax.typing import ArrayLike
@@ -562,6 +563,46 @@ def f(a):
         _assert_pytree_isequal(result, result_ref)
 
 
+@pytest.mark.parametrize("use_jit", [True, False])
+def test_tesseract_loss(served_vectoradd_tesseract, use_jit):
+    vectoradd_tess = Tesseract(served_vectoradd_tesseract)
+    a = np.array([1.0, 2.0, 3.0], dtype="float32")
+
+    # b = jax.lax.stop_gradient(b)
+
+    def loss_fn(a):
+        b = np.array([4.0, 5.0, 6.0], dtype="float32")
+
+        vectoradd_fn_a: jax.Callable = lambda a: apply_tesseract(
+            vectoradd_tess,
+            inputs=dict(
+                a=a,
+                b=b,
+            ),
+        )
+
+        c = vectoradd_fn_a(a)["c"]
+
+        vectoradd_fn_b: jax.Callable = lambda a: apply_tesseract(
+            vectoradd_tess,
+            inputs=dict(
+                a=a,
+                b=c,
+            ),
+        )
+
+        outputs = vectoradd_fn_b(a)
+
+        return jnp.sum((outputs["c"]) ** 2)
+
+    if use_jit:
+        loss_fn = jax.jit(loss_fn)
+
+    value_and_grad_fn = jax.value_and_grad(loss_fn)
+
+    assert value_and_grad_fn(a) is not None
+
+
 def test_non_abstract_tesseract_vjp(served_non_abstract_tesseract):
     non_abstract_tess = Tesseract(served_non_abstract_tesseract)
 

tests/vectoradd_tesseract/tesseract_api.py

@@ -0,0 +1,44 @@
+# Copyright 2025 Pasteur Labs. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+from tesseract_core.runtime import Array, Differentiable, Float32
+
+
+class InputSchema(BaseModel):
+    a: Differentiable[Array[(None,), Float32]] = Field(description="Arbitrary vector a")
+    b: Array[(None,), Float32] = Field(description="Arbitrary vector b")
+
+
+class OutputSchema(BaseModel):
+    c: Differentiable[Array[(None,), Float32]] = Field(
+        description="Vector s_a·a + s_b·b"
+    )
+
+
+def apply(inputs: InputSchema) -> OutputSchema:
+    """Adds two vectors `a` and `b`."""
+    return OutputSchema(
+        c=inputs.a + inputs.b,
+    )
+
+
+def abstract_eval(abstract_inputs):
+    """Abstract evaluation of the addition operation."""
+    return {
+        "c": abstract_inputs.a,
+    }
+
+
+def vector_jacobian_product(
+    inputs: InputSchema,
+    vjp_inputs: set[str],
+    vjp_outputs: set[str],
+    cotangent_vector: dict[str, Any],
+):
+    return {
+        "a": cotangent_vector["c"],
+    }

tests/vectoradd_tesseract/tesseract_config.yaml

@@ -0,0 +1,9 @@
+name: vectoradd_tesseracts
+version: "2025-11-05"
+description: |
+  Tesseract that adds two vectors. Uses jax internally.
+
+build_config:
+  target_platform: "native"
+  # package_data: []
+  # custom_build_steps: []

tests/vectoradd_tesseract/tesseract_requirements.txt

@@ -0,0 +1 @@
+jax[cpu]