doc: Fix RBF fitting example (#25)

#### Relevant issue or PR
- #3

#### Description of changes
- [x] Python API (i.e. `tesseract_core.sdk.tesseract`) now used to
launch and interact with Tesseract
- [x] Jacobian argument order corrected
- [x] Only `'weights'` is provided in `diff_inputs` now (`length_scale`
not used)
- [x] Offered options for using either jacobian or vjp
- [x] Pydantic model validators now return `self` instead of `None`
- [x] Removed port argument
- [x] Documentation updated
- [x] Tidied up imports and namespaces for jax recipe and
`vectoradd_jax` unit tesseract

#### Testing done
- [ ] CI passes
- [x] Tesseract builds
- [x] `optimization_routine.py` runs successfully with graphs 

#### License

- [x] By submitting this pull request, I confirm that my contribution is
made under the terms of the [Apache 2.0
license](https://pasteurlabs.github.io/tesseract/LICENSE).
- [x] I sign the Developer Certificate of Origin below by adding my name
and email address to the `Signed-off-by` line.

<details>
<summary><b>Developer Certificate of Origin</b></summary>

```text
Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.


Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.
```

</details>

Signed-off-by: Jonathan Brodrick <jonathan.brodrick@simulation.science>

---------

Co-authored-by: Matteo Salvador <44835004+MatteoSalvador@users.noreply.github.com>
Co-authored-by: Dion Häfner <dion.haefner@simulation.science>

Author: 3 people

examples/sciml/rbf_fitting/optimization_routine.py

@@ -3,25 +3,13 @@
 
 """This is an optimisation scripts for the RBF fitting example."""
 
-import argparse
-
 import jax.numpy as jnp
+import matplotlib.pyplot as plt
 import optax
 from jax import random
-from matplotlib import pyplot as plt
 from plotting import plot_loss, plot_rbf
-from tesseract_client import Client
-
-parser = argparse.ArgumentParser(description="RBF tesseract optimisation script")
-parser.add_argument(
-    "-p", "--port", help="Port at which RBF tesseract is being served", required=True
-)
-args = vars(parser.parse_args())
 
-if "port" not in args:
-    raise ValueError("Port at which RBF tesseract is being served is required.")
-
-port = int(args["port"])
+from tesseract_core.sdk.tesseract import Tesseract
 
 # JAX random keys
 key = random.PRNGKey(42)
@@ -54,11 +42,8 @@ def ground_truth(x):
     "y_target": y_target.tolist(),
 }
 
-jac_inputs = ["weights", "length_scale"]
-jac_outputs = ["mse"]
-
-# Initialize Tesseract client
-client = Client(host="127.0.0.1", port=port)
+diff_inputs = ["weights"]
+diff_outputs = ["mse"]
 
 # Initialize optimizer
 optimizer = optax.adam(learning_rate=0.5)
@@ -68,34 +53,35 @@ def ground_truth(x):
 loss_history = []
 max_iterations = 100
 print(f"Starting the optimization process for {max_iterations} iterations.")
-for n_iteration in range(max_iterations):
-    if n_iteration % 10 == 0:
-        print(f" ---- iteration {n_iteration} / {max_iterations}")
-
-    # Compute loss
-    apply_response = client.request("apply", method="POST", payload={"inputs": inputs})
-    loss = apply_response["mse"]["data"]["buffer"]
-    loss_history.append(loss)
-
-    # Compute gradients
-    jacobian_response = client.request(
-        "jacobian",
-        method="POST",
-        payload={
-            "inputs": inputs,
-            "jac_inputs": jac_inputs,
-            "jac_outputs": jac_outputs,
-        },
-    )
 
-    buffer = jacobian_response["weights"]["mse"]["data"]["buffer"]
-    grad_weights = jnp.array(buffer, dtype=jnp.float32)
-
-    # Update weights
-    weights = jnp.array(inputs["weights"])
-    updates, opt_state = optimizer.update(grad_weights, opt_state, weights)
-    weights = optax.apply_updates(weights, updates)
-    inputs["weights"] = weights.tolist()
+# Initialize Tesseract client
+with Tesseract.from_image(image="rbf_fitting") as tess:
+    for n_iteration in range(max_iterations):
+        if n_iteration % 10 == 0:
+            print(f" ---- iteration {n_iteration} / {max_iterations}")
+
+        # Compute loss
+        apply_response = tess.apply(inputs)
+        loss = apply_response["mse"]
+        loss_history.append(loss)
+
+        # Compute gradients
+        # Option 1: Use Jacobian
+        # grad_weights = tess.jacobian(
+        #     inputs,
+        #     diff_inputs,
+        #     diff_outputs,
+        # )["mse"]["weights"]
+        # Option 2: Use VJP
+        grad_weights = tess.vector_jacobian_product(
+            inputs, diff_inputs, diff_outputs, {"mse": 1.0}
+        )["weights"]
+
+        # Update weights
+        weights = jnp.array(inputs["weights"])
+        updates, opt_state = optimizer.update(grad_weights, opt_state, weights)
+        weights = optax.apply_updates(weights, updates)
+        inputs["weights"] = weights.tolist()
 
 print("Optimisation completed!")
 

examples/sciml/rbf_fitting/tesseract_api.py

@@ -1,11 +1,17 @@
 # Copyright 2025 Pasteur Labs. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import jax
+from functools import partial
+from typing import Any
+
 import jax.numpy as jnp
+import jax.tree
+from jax import ShapeDtypeStruct, eval_shape, jacrev, jit, jvp, vjp
 from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
 
 from tesseract_core.runtime import Array, Differentiable, Float32
+from tesseract_core.runtime.tree_transforms import filter_func, flatten_with_paths
 
 
 def gaussian_rbf(x: float, c: float, length_scale: float) -> float:
@@ -25,7 +31,13 @@ def mse_error(
     for coeff, center in zip(weights, x_centers, strict=False):
         y_hat += coeff * gaussian_rbf(x_target, center, length_scale)
 
-    return jnp.mean((y_target - y_hat) ** 2)
+    return {"mse": jnp.mean((y_target - y_hat) ** 2)}
+
+
+@jit
+def apply_jit(inputs: dict) -> dict:
+    ordered_keys = ["x_centers", "weights", "length_scale", "x_target", "y_target"]
+    return mse_error(*(inputs[key] for key in ordered_keys))
 
 
 #
@@ -49,20 +61,22 @@ class InputSchema(BaseModel):
     )
 
     @model_validator(mode="after")
-    def validate_shape_targets(self) -> None:
+    def validate_shape_targets(self) -> Self:
         if self.x_target.shape != self.y_target.shape:
             raise ValueError(
                 f"x_target and y_target must have the same shape. "
                 f"Got {self.x_target.shape} and {self.y_target.shape} instead."
             )
+        return self
 
     @model_validator(mode="after")
-    def validate_shape_weights(self) -> None:
+    def validate_shape_weights(self) -> Self:
         if self.x_centers.shape != self.weights.shape:
             raise ValueError(
                 f"x_centers and weights must have the same shape. "
                 f"Got {self.x_centers.shape} and {self.weights.shape} instead."
             )
+        return self
 
 
 class OutputSchema(BaseModel):
@@ -77,18 +91,11 @@ class OutputSchema(BaseModel):
 
 
 def apply(inputs: InputSchema) -> OutputSchema:
-    mse = mse_error(
-        inputs.x_centers,
-        inputs.weights,
-        inputs.length_scale,
-        inputs.x_target,
-        inputs.y_target,
-    )
-    return OutputSchema(mse=mse)
+    return apply_jit(inputs.model_dump())
 
 
 #
-# Optional endpoints
+# Jax-handled AD endpoints (no need to modify)
 #
 
 
@@ -97,20 +104,86 @@ def jacobian(
     jac_inputs: set[str],
     jac_outputs: set[str],
 ):
-    mse_signature = ["x_centers", "weights", "length_scale", "x_target", "y_target"]
-
-    jac_result = {}
-    for dx in jac_inputs:
-        grad_func = jax.jacrev(mse_error, argnums=mse_signature.index(dx))
-        for dy in jac_outputs:
-            jac_result[dx] = {
-                dy: grad_func(
-                    inputs.x_centers,
-                    inputs.weights,
-                    inputs.length_scale,
-                    inputs.x_target,
-                    inputs.y_target,
-                )
-            }
-
-    return jac_result
+    return jac_jit(inputs.model_dump(), tuple(jac_inputs), tuple(jac_outputs))
+
+
+def jacobian_vector_product(
+    inputs: InputSchema,
+    jvp_inputs: set[str],
+    jvp_outputs: set[str],
+    tangent_vector: dict[str, Any],
+):
+    return jvp_jit(
+        inputs.model_dump(),
+        tuple(jvp_inputs),
+        tuple(jvp_outputs),
+        tangent_vector,
+    )
+
+
+def vector_jacobian_product(
+    inputs: InputSchema,
+    vjp_inputs: set[str],
+    vjp_outputs: set[str],
+    cotangent_vector: dict[str, Any],
+):
+    return vjp_jit(
+        inputs.model_dump(),
+        tuple(vjp_inputs),
+        tuple(vjp_outputs),
+        cotangent_vector,
+    )
+
+
+def abstract_eval(abstract_inputs):
+    """Calculate output shape of apply from the shape of its inputs."""
+    jaxified_inputs = jax.tree.map(
+        lambda x: ShapeDtypeStruct(**x),
+        abstract_inputs.model_dump(),
+        is_leaf=lambda x: (x.keys() == {"shape", "dtype"}),
+    )
+    jax_shapes = eval_shape(apply_jit, jaxified_inputs)
+    return jax.tree.map(
+        lambda sd: {"shape": sd.shape, "dtype": str(sd.dtype)}, jax_shapes
+    )
+
+
+#
+# Helper functions
+#
+
+
+@partial(jit, static_argnames=["jac_inputs", "jac_outputs"])
+def jac_jit(
+    inputs: dict,
+    jac_inputs: tuple[str],
+    jac_outputs: tuple[str],
+):
+    filtered_apply = filter_func(apply_jit, inputs, jac_outputs)
+    return jacrev(filtered_apply)(flatten_with_paths(inputs, include_paths=jac_inputs))
+
+
+@partial(jit, static_argnames=["jvp_inputs", "jvp_outputs"])
+def jvp_jit(
+    inputs: dict, jvp_inputs: tuple[str], jvp_outputs: tuple[str], tangent_vector: dict
+):
+    filtered_apply = filter_func(apply_jit, inputs, jvp_outputs)
+    return jvp(
+        filtered_apply,
+        [flatten_with_paths(inputs, include_paths=jvp_inputs)],
+        [tangent_vector],
+    )[1]
+
+
+@partial(jit, static_argnames=["vjp_inputs", "vjp_outputs"])
+def vjp_jit(
+    inputs: dict,
+    vjp_inputs: tuple[str],
+    vjp_outputs: tuple[str],
+    cotangent_vector: dict,
+):
+    filtered_apply = filter_func(apply_jit, inputs, vjp_outputs)
+    _, vjp_func = vjp(
+        filtered_apply, flatten_with_paths(inputs, include_paths=vjp_inputs)
+    )
+    return vjp_func(cotangent_vector)[0]

examples/unit_tesseracts/vectoradd_jax/tesseract_api.py

@@ -6,7 +6,7 @@
 
 import jax.numpy as jnp
 import jax.tree
-from jax import ShapeDtypeStruct, eval_shape, jit, jvp, vjp
+from jax import ShapeDtypeStruct, eval_shape, jacrev, jit, jvp, vjp
 from pydantic import BaseModel, Field, model_validator
 from typing_extensions import Self
 
@@ -164,6 +164,4 @@ def jac_jit(
     jac_outputs: tuple[str],
 ):
     filtered_apply = filter_func(apply_jit, inputs, jac_outputs)
-    return jax.jacrev(filtered_apply)(
-        flatten_with_paths(inputs, include_paths=jac_inputs)
-    )
+    return jacrev(filtered_apply)(flatten_with_paths(inputs, include_paths=jac_inputs))

tesseract_core/sdk/templates/jax/tesseract_api.py

@@ -8,7 +8,7 @@
 from typing import Any
 
 import jax.tree
-from jax import ShapeDtypeStruct, eval_shape, jit, jvp, vjp
+from jax import ShapeDtypeStruct, eval_shape, jacrev, jit, jvp, vjp
 from pydantic import BaseModel
 
 from tesseract_core.runtime import Differentiable, Float32
@@ -123,9 +123,7 @@ def jac_jit(
     jac_outputs: tuple[str],
 ):
     filtered_apply = filter_func(apply_jit, inputs, jac_outputs)
-    return jax.jacrev(filtered_apply)(
-        flatten_with_paths(inputs, include_paths=jac_inputs)
-    )
+    return jacrev(filtered_apply)(flatten_with_paths(inputs, include_paths=jac_inputs))
 
 
 @partial(jit, static_argnames=["jvp_inputs", "jvp_outputs"])