doc: more tweaks based on beta feedback (#15)

#### Relevant issue or PR
n/a

#### Description of changes
see files tab

#### Testing done
:eyes:

#### 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

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:

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:

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).