Merge branch 'main' into jacan/demos

Author: jacanchaplais

README.md

@@ -2,7 +2,7 @@
 
 ### Tesseract-JAX
 
-`tesseract-jax` executes [Tesseracts](https://github.com/pasteurlabs/tesseract-core) as part of [JAX](https://github.com/jax-ml/jax) programs, with full support for function transformations like JIT, `grad`, and more.
+Tesseract-JAX is a lightweight extension to [Tesseract Core](https://github.com/pasteurlabs/tesseract-core) that makes Tesseracts look and feel like regular [JAX](https://github.com/jax-ml/jax) primitives, and makes them jittable, differentiable, and composable.
 
 [Read the docs](https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/) |
 [Explore the examples](https://github.com/pasteurlabs/tesseract-jax/tree/main/examples) |
@@ -12,7 +12,16 @@
 
 ---
 
-The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesseract_client, inputs)`](https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/content/api.html#tesseract_jax.apply_tesseract), which is fully traceable by JAX. This enables end-to-end autodifferentiation and JIT compilation of Tesseract-based pipelines.
+The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesseract_client, inputs)`](https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/content/api.html#tesseract_jax.apply_tesseract), which is fully traceable by JAX. This enables end-to-end autodifferentiation and JIT compilation of Tesseract-based pipelines:
+
+```python
+@jax.jit
+def vector_sum(x, y):
+    res = apply_tesseract(vectoradd_tesseract, {"a": {"v": x}, "b": {"v": y}})
+    return res["vector_add"]["result"].sum()
+
+jax.grad(vector_sum)(x, y) # 🎉
+```
 
 ## Quick start
 
@@ -31,7 +40,8 @@ The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesser
 2. Build an example Tesseract:
 
    ```bash
-   $ tesseract build examples/simple/vectoradd_jax
+   $ git clone https://github.com/pasteurlabs/tesseract-jax
+   $ tesseract build tesseract-jax/examples/simple/vectoradd_jax
    ```
 
 3. Use it as part of a JAX program via the JAX-native `apply_tesseract` function:
@@ -62,22 +72,23 @@ The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesser
 
    vector_sum_grad = jax.grad(vector_sum)
    vector_sum_grad(x, y)
-    ```
+   ```
 
 > [!TIP]
 > Now you're ready to jump into our [examples](https://github.com/pasteurlabs/tesseract-jax/tree/main/examples) for more ways to use Tesseract-JAX.
 
 ## Sharp edges
 
-- **Arrays vs. array-like objects**: Tesseract-JAX ist stricter than Tesseract Core in that all array inputs to Tesseracts must be JAX or NumPy arrays, not just any array-like (such as Python floats or lists). As a result, you may need to convert your inputs to JAX arrays before passing them to Tesseract-JAX, including scalar values.
+- **Arrays vs. array-like objects**: Tesseract-JAX is stricter than Tesseract Core in that all array inputs to Tesseracts must be JAX or NumPy arrays, not just any array-like (such as Python floats or lists). As a result, you may need to convert your inputs to JAX arrays before passing them to Tesseract-JAX, including scalar values.
 
   ```python
   from tesseract_core import Tesseract
   from tesseract_jax import apply_tesseract
 
-  tess = Tesseract.from_image("vectoradd")
-  apply_tesseract(tess, {"a": 1.0, "b": 2.0})  # ❌ raises an error
-  apply_tesseract(tess, {"a": jnp.array(1.0), "b": jnp.array(2.0)})  # ✅ works
+  tess = Tesseract.from_image("vectoradd_jax")
+  with Tesseract.from_image("vectoradd_jax") as tess:
+      apply_tesseract(tess, {"a": {"v": [1.0]}, "b": {"v": [2.0]}})  # ❌ raises an error
+      apply_tesseract(tess, {"a": {"v": jnp.array([1.0])}, "b": {"v": jnp.array([2.0])}})  # ✅ works
   ```
 - **Additional required endpoints**: Tesseract-JAX requires the [`abstract_eval`](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/api/endpoints.html#abstract-eval) Tesseract endpoint to be defined for all operations. This is because JAX mandates abstract evaluation of all operations before they are executed. Additionally, many gradient transformations like `jax.grad` require [`vector_jacobian_product`](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/api/endpoints.html#vector-jacobian-product) to be defined.
 

docs/content/get-started.md

@@ -1,9 +1,5 @@
 # Get started
 
-`tesseract-jax` executes [Tesseracts](https://github.com/pasteurlabs/tesseract-core) as part of JAX programs, with full support for function transformations like JIT, `grad`, `jvp`, and more.
-
-The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesseract_client, inputs)`](tesseract_jax.apply_tesseract), which is fully traceable by JAX. This enables end-to-end autodifferentiation and JIT compilation of Tesseract-based pipelines.
-
 ## Quick start
 
 ```{note}
@@ -23,7 +19,8 @@ For more detailed installation instructions, please refer to the [Tesseract Core
 2. Build an example Tesseract:
 
    ```bash
-   $ tesseract build examples/simple/vectoradd_jax
+   $ git clone https://github.com/pasteurlabs/tesseract-jax
+   $ tesseract build tesseract-jax/examples/simple/vectoradd_jax
    ```
 
 3. Use it as part of a JAX program:
@@ -62,15 +59,16 @@ Now you're ready to jump into our [examples](https://github.com/pasteurlabs/tess
 
 ## Sharp edges
 
-- **Arrays vs. array-like objects**: Tesseract-JAX ist stricter than Tesseract Core in that all array inputs to Tesseracts must be JAX or NumPy arrays, not just any array-like (such as Python floats or lists). As a result, you may need to convert your inputs to JAX arrays before passing them to Tesseract-JAX, including scalar values.
+- **Arrays vs. array-like objects**: Tesseract-JAX is stricter than Tesseract Core in that all array inputs to Tesseracts must be JAX or NumPy arrays, not just any array-like (such as Python floats or lists). As a result, you may need to convert your inputs to JAX arrays before passing them to Tesseract-JAX, including scalar values.
 
   ```python
   from tesseract_core import Tesseract
   from tesseract_jax import apply_tesseract
 
-  tess = Tesseract.from_image("vectoradd")
-  apply_tesseract(tess, {"a": 1.0, "b": 2.0})  # ❌ raises an error
-  apply_tesseract(tess, {"a": jnp.array(1.0), "b": jnp.array(2.0)})  # ✅ works
+  tess = Tesseract.from_image("vectoradd_jax")
+  with Tesseract.from_image("vectoradd_jax") as tess:
+      apply_tesseract(tess, {"a": {"v": [1.0]}, "b": {"v": [2.0]}})  # ❌ raises an error
+      apply_tesseract(tess, {"a": {"v": jnp.array([1.0])}, "b": {"v": jnp.array([2.0])}})  # ✅ works
   ```
 - **Additional required endpoints**: Tesseract-JAX requires the [`abstract_eval`](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/api/endpoints.html#abstract-eval) Tesseract endpoint to be defined for all operations. This is because JAX mandates abstract evaluation of all operations before they are executed. Additionally, many gradient transformations like `jax.grad` require [`vector_jacobian_product`](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/api/endpoints.html#vector-jacobian-product) to be defined.
 

docs/index.md

@@ -1,9 +1,20 @@
 # Tesseract-JAX
 
-```{include} content/get-started.md
-:start-line: 2
+Tesseract-JAX is a lightweight extension to [Tesseract Core](https://github.com/pasteurlabs/tesseract-core) that makes Tesseracts look and feel like regular [JAX](https://github.com/jax-ml/jax) primitives, and makes them jittable, differentiable, and composable.
+
+The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesseract_client, inputs)`](tesseract_jax.apply_tesseract), which is fully traceable by JAX. This enables end-to-end autodifferentiation and JIT compilation of Tesseract-based pipelines:
+
+```python
+@jax.jit
+def vector_sum(x, y):
+    res = apply_tesseract(vectoradd_tesseract, {"a": {"v": x}, "b": {"v": y}})
+    return res["vector_add"]["result"].sum()
+
+jax.grad(vector_sum)(x, y) # 🎉
 ```
 
+Want to learn more? See how to [get started](content/get-started.md) with Tesseract-JAX, explore the [API reference](content/api.md), or learn by [example](demo_notebooks/simple.ipynb).
+
 ## License
 
 Tesseract JAX is licensed under the [Apache License 2.0](https://github.com/pasteurlabs/tesseract-jax/LICENSE) and is free to use, modify, and distribute (under the terms of the license).

examples/simple/vectoradd_jax/tesseract_api.py

@@ -65,16 +65,20 @@ def apply_jit(inputs: dict) -> dict:
     b_scaled = inputs["b"]["s"] * inputs["b"]["v"]
     add_result = a_scaled + b_scaled
     min_result = a_scaled - b_scaled
+
+    def safe_norm(x, ord):
+        # Compute the norm of a vector, adding a small epsilon to ensure
+        # differentiability and avoid division by zero
+        return jnp.power(jnp.power(jnp.abs(x), ord).sum() + 1e-8, 1 / ord)
+
     return {
         "vector_add": {
             "result": add_result,
-            "normed_result": add_result
-            / jnp.linalg.norm(add_result, ord=inputs["norm_ord"]),
+            "normed_result": add_result / safe_norm(add_result, ord=inputs["norm_ord"]),
         },
         "vector_min": {
             "result": min_result,
-            "normed_result": min_result
-            / jnp.linalg.norm(min_result, ord=inputs["norm_ord"]),
+            "normed_result": min_result / safe_norm(min_result, ord=inputs["norm_ord"]),
         },
     }