doc: add rendered demos to docs (#9)

#### Relevant issue or PR
n/a

#### Description of changes
- Add rendered demos to docs.
- Add extra navigation item for index page.

#### Testing done
manual

#### 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-jax/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: Dion Häfner <[email protected]>

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](https://github.com/jax-ml/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

@@ -7,12 +7,14 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
 import re
+import shutil
+from pathlib import Path
 
 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__)
@@ -28,12 +30,16 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
 extensions = [
-    "myst_parser",
+    "myst_nb",  # This is myst-parser + jupyter notebook support
     "sphinx.ext.intersphinx",
     "sphinx.ext.autodoc",
     "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 = [
@@ -62,6 +68,18 @@
 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"]
+
+
+# -- Handle Jupyter notebooks ------------------------------------------------
+
+# Do not execute notebooks during build (just take existing output)
+nb_execution_mode = "off"
+
+# Copy example notebooks to demo_notebooks folder on every build
+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 @@
-# API documentation
+# API reference
 
 ```{eval-rst}
 .. automodule:: tesseract_jax

docs/conf.py

@@ -0,0 +1,79 @@
+# 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}
+Before proceeding, make sure you have a [working installation of Docker](https://docs.docker.com/engine/install/) and a modern Python installation (Python 3.10+).
+```
+
+```{seealso}
+For more detailed installation instructions, please refer to the [Tesseract Core documentation](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/introduction/installation.html).
+```
+
+1. Install Tesseract-JAX:
+
+   ```bash
+   $ pip install tesseract-jax
+   ```
+
+2. Build an example Tesseract:
+
+   ```bash
+   $ tesseract build examples/simple/vectoradd_jax
+   ```
+
+3. Use it as part of a JAX program:
+
+   ```python
+   import jax
+   import jax.numpy as jnp
+   from tesseract_core import Tesseract
+   from tesseract_jax import apply_tesseract
+
+   # Load the Tesseract
+   t = Tesseract.from_image("vectoradd_jax")
+   t.serve()
+
+   # Run it with JAX
+   x = jnp.ones((1000,))
+   y = jnp.ones((1000,))
+
+   def vector_sum(x, y):
+       res = apply_tesseract(t, {"a": {"v": x}, "b": {"v": y}})
+       return res["vector_add"]["result"].sum()
+
+   vector_sum(x, y) # success!
+
+   # 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_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 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.
+
+  ```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
+  ```
+- **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.
+
+```{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`.
+```

docs/content/api.md

@@ -0,0 +1 @@
+*.ipynb

docs/content/get-started.md

@@ -1,93 +1,39 @@
-# tesseract-jax
+# Tesseract-JAX
 
-`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}
-Before proceeding, make sure you have a [working installation of Docker](https://docs.docker.com/engine/install/) and a modern Python installation (Python 3.10+).
-```
-
-```{seealso}
-For more detailed installation instructions, please refer to the [Tesseract Core documentation](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/introduction/installation.html).
+```{include} content/get-started.md
+:start-line: 2
 ```
 
-1. Install Tesseract-JAX:
-
-   ```bash
-   $ pip install tesseract-jax
-   ```
-
-2. Build an example Tesseract:
-
-   ```bash
-   $ tesseract build demo/simple/vectoradd_jax
-   ```
-
-3. Use it as part of a JAX program:
-
-   ```python
-   import jax
-   import jax.numpy as jnp
-   from tesseract_core import Tesseract
-   from tesseract_jax import apply_tesseract
-
-   # Load the Tesseract
-   t = Tesseract.from_image("vectoradd_jax")
-
-   # Run it with JAX
-   x = jnp.ones((1000, 1000))
-   y = jnp.ones((1000, 1000))
+## License
 
-   def vector_add(x, y):
-       return apply_tesseract(t, x, y)
+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).
 
-    vector_add(x, y) # success!
+Tesseract is a registered trademark of Pasteur Labs, Inc. and may not be used without permission.
 
-    # You can also use it with JAX transformations like JIT and grad
-    vector_add_jit = jax.jit(vector_add)
-    vector_add_jit(x, y)
 
-    vector_add_grad = jax.grad(vector_add)
-    vector_add_grad(x, y)
-    ```
+```{toctree}
+:caption: Usage
+:maxdepth: 2
+:hidden:
 
-```{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.
+content/get-started
+content/api
 ```
 
-## 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.
-
-  ```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
-  ```
-- **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.
+```{toctree}
+:caption: Examples
+:maxdepth: 2
+:hidden:
 
-```{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`.
+demo_notebooks/simple.ipynb
+demo_notebooks/cfd.ipynb
 ```
 
 ```{toctree}
-:caption: Contents
+:caption: See also
 :maxdepth: 2
 :hidden:
 
-content/api
-Tesseract Core documentation <https://docs.pasteurlabs.ai/projects/tesseract-core/latest/>
+Tesseract Core docs <https://docs.pasteurlabs.ai/projects/tesseract-core/latest/>
 Tesseract User Forums <https://si-tesseract.discourse.group/>
 ```
-
-## 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.