overhaul demos (now examples) and docs

Author: dionhaefner

.github/workflows/test_examples.yml

@@ -16,7 +16,7 @@ jobs:
         # test with oldest supported Python version only (for slow tests)
         python-version: ["3.10"]
 
-        demo:
+        example:
           - simple
           - cfd
 
@@ -46,7 +46,7 @@ jobs:
           uv sync --extra dev --frozen
 
       - name: Run example
-        working-directory: demo/${{matrix.demo}}
+        working-directory: examples/${{matrix.example}}
         run: |
           uv pip install jupyter
           uv run --no-sync jupyter nbconvert --to notebook --execute demo.ipynb

README.md

@@ -5,7 +5,7 @@
 `tesseract-jax` executes [Tesseracts](https://github.com/pasteurlabs/tesseract-core) as part of JAX programs, with full support for function transformations like JIT, `grad`, and more.
 
 [Read the docs](https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/) |
-[Explore the demos](https://github.com/pasteurlabs/tesseract-jax/tree/main/demo) |
+[Explore the examples](https://github.com/pasteurlabs/tesseract-jax/tree/main/examples) |
 [Report an issue](https://github.com/pasteurlabs/tesseract-jax/issues) |
 [Talk to the community](https://si-tesseract.discourse.group/) |
 [Contribute](CONTRIBUTING.md)
@@ -31,7 +31,7 @@ The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesser
 2. Build an example Tesseract:
 
    ```bash
-   $ tesseract build demo/simple/vectoradd_jax
+   $ tesseract build examples/simple/vectoradd_jax
    ```
 
 3. Use it as part of a JAX program via the JAX-native `apply_tesseract` function:
@@ -44,26 +44,28 @@ The API of Tesseract-JAX consists of a single function, [`apply_tesseract(tesser
 
    # Load the Tesseract
    t = Tesseract.from_image("vectoradd_jax")
+   t.serve()
 
    # Run it with JAX
-   x = jnp.ones((1000, 1000))
-   y = jnp.ones((1000, 1000))
+   x = jnp.ones((1000,))
+   y = jnp.ones((1000,))
 
-   def vector_add(x, y):
-       return apply_tesseract(t, x, y)
+   def vector_sum(x, y):
+       res = apply_tesseract(t, {"a": {"v": x}, "b": {"v": y}})
+       return res["vector_add"]["result"].sum()
 
-    vector_add(x, y) # success!
+   vector_sum(x, y) # success!
 
-    # You can also use it with JAX transformations like JIT and grad
-    vector_add_jit = jax.jit(vector_add)
-    vector_add_jit(x, y)
+   # You can also use it with JAX transformations like JIT and grad
+   vector_sum_jit = jax.jit(vector_sum)
+   vector_sum_jit(x, y)
 
-    vector_add_grad = jax.grad(vector_add)
-    vector_add_grad(x, y)
+   vector_sum_grad = jax.grad(vector_sum)
+   vector_sum_grad(x, y)
     ```
 
 > [!TIP]
-> Now you're ready to jump into our [demos](https://github.com/pasteurlabs/tesseract-jax/tree/main/demo) for more examples on how to use Tesseract-JAX.
+> 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
 

demo/cfd/Readme.md

@@ -12,9 +12,9 @@
 
 from tesseract_jax import __version__
 
-project = "tesseract-jax"
-copyright = "2025, The tesseract-jax Team @ Pasteur Labs"
-author = "The tesseract-jax Team @ Pasteur Labs"
+project = "Tesseract-JAX"
+copyright = "2025, Pasteur Labs"
+author = "The Tesseract-JAX Team @ Pasteur Labs + OSS contributors"
 
 # The short X.Y version
 parsed_version = re.match(r"(\d+\.\d+\.\d+)", __version__)
@@ -36,6 +36,10 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
     "sphinx_autodoc_typehints",
+    # Copy button for code blocks
+    "sphinx_copybutton",
+    # OpenGraph metadata for social media sharing
+    "sphinxext.opengraph",
 ]
 
 myst_enable_extensions = [
@@ -64,7 +68,7 @@
 html_theme_options = {
     "light_logo": "logo-light.png",
     "dark_logo": "logo-dark.png",
-    "sidebar_hide_name": False,
+    "sidebar_hide_name": True,
 }
 html_css_files = ["custom.css"]
 
@@ -75,7 +79,7 @@
 nb_execution_mode = "off"
 
 # Copy example notebooks to demo_notebooks folder on every build
-for example_notebook in Path("../demo").glob("*/demo.ipynb"):
+for example_notebook in Path("../examples").glob("*/demo.ipynb"):
     # Copy the example notebook to the docs folder
     dest = (Path("demo_notebooks") / example_notebook.parent.name).with_suffix(".ipynb")
     shutil.copyfile(example_notebook, dest)

demo/simple/README.md

@@ -1,4 +1,4 @@
-# Tesseract-JAX
+# 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.
 
@@ -23,7 +23,7 @@ For more detailed installation instructions, please refer to the [Tesseract Core
 2. Build an example Tesseract:
 
    ```bash
-   $ tesseract build demo/simple/vectoradd_jax
+   $ tesseract build examples/simple/vectoradd_jax
    ```
 
 3. Use it as part of a JAX program:
@@ -36,26 +36,28 @@ For more detailed installation instructions, please refer to the [Tesseract Core
 
    # Load the Tesseract
    t = Tesseract.from_image("vectoradd_jax")
+   t.serve()
 
    # Run it with JAX
-   x = jnp.ones((1000, 1000))
-   y = jnp.ones((1000, 1000))
+   x = jnp.ones((1000,))
+   y = jnp.ones((1000,))
 
-   def vector_add(x, y):
-       return apply_tesseract(t, x, y)
+   def vector_sum(x, y):
+       res = apply_tesseract(t, {"a": {"v": x}, "b": {"v": y}})
+       return res["vector_add"]["result"].sum()
 
-    vector_add(x, y) # success!
+   vector_sum(x, y) # success!
 
-    # You can also use it with JAX transformations like JIT and grad
-    vector_add_jit = jax.jit(vector_add)
-    vector_add_jit(x, y)
+   # You can also use it with JAX transformations like JIT and grad
+   vector_sum_jit = jax.jit(vector_sum)
+   vector_sum_jit(x, y)
 
-    vector_add_grad = jax.grad(vector_add)
-    vector_add_grad(x, y)
-    ```
+   vector_sum_grad = jax.grad(vector_sum)
+   vector_sum_grad(x, y)
+   ```
 
 ```{tip}
-Now you're ready to jump into our [demos](https://github.com/pasteurlabs/tesseract-jax/tree/main/demo) for more examples on how to use Tesseract-JAX.
+Now you're ready to jump into our [examples](https://github.com/pasteurlabs/tesseract-jax/tree/main/examples) for ways to use Tesseract-JAX.
 ```
 
 ## Sharp edges
@@ -75,9 +77,3 @@ Now you're ready to jump into our [demos](https://github.com/pasteurlabs/tessera
 ```{tip}
 When creating a new Tesseract based on a JAX function, use `tesseract init --recipe jax` to define all required endpoints automatically, including `abstract_eval` and `vector_jacobian_product`.
 ```
-
-## 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).
-
-Tesseract is a registered trademark of Pasteur Labs, Inc. and may not be used without permission.

docs/conf.py

@@ -1,23 +1,39 @@
-```{include} content/index.md
+# Tesseract-JAX
+
+```{include} content/get-started.md
+:start-line: 2
 ```
 
+## 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).
+
+Tesseract is a registered trademark of Pasteur Labs, Inc. and may not be used without permission.
+
 
 ```{toctree}
-:caption: Contents
+:caption: Usage
 :maxdepth: 2
 :hidden:
 
-content/index
+content/get-started
 content/api
-Tesseract Core docs <https://docs.pasteurlabs.ai/projects/tesseract-core/latest/>
-Tesseract User Forums <https://si-tesseract.discourse.group/>
 ```
 
 ```{toctree}
-:caption: Demos
+:caption: Examples
 :maxdepth: 2
 :hidden:
 
 demo_notebooks/simple.ipynb
 demo_notebooks/cfd.ipynb
 ```
+
+```{toctree}
+:caption: See also
+:maxdepth: 2
+:hidden:
+
+Tesseract Core docs <https://docs.pasteurlabs.ai/projects/tesseract-core/latest/>
+Tesseract User Forums <https://si-tesseract.discourse.group/>
+```

docs/content/index.md renamed to docs/content/get-started.md

@@ -0,0 +1,8 @@
+# Tesseract-JAX examples
+
+This directory contains example Tesseract configurations, notebooks. and scripts demonstrating how to use Tesseract-JAX in various contexts. Each example is self-contained and can be run independently.
+
+## Examples
+
+- [Simple](simple/demo.ipynb): A basic example of using Tesseract-JAX with a simple vector addition task. It demonstrates how to build a Tesseract and execute it within JAX.
+- [CFD](cfd/demo.ipynb): A more complex example demonstrating how to use Tesseract-JAX to differentiate through a computational fluid dynamics (CFD) simulation Tesseract.