From 6750f231f06059de5423ed027c9711d65f5ad10e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dion=20H=C3=A4fner?= Date: Wed, 23 Apr 2025 20:57:54 +0200 Subject: [PATCH 1/3] dry --- docs/content/get-started.md | 4 ---- docs/index.md | 8 +++++--- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/docs/content/get-started.md b/docs/content/get-started.md index 0232fc6..156668f 100644 --- a/docs/content/get-started.md +++ b/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} diff --git a/docs/index.md b/docs/index.md index fa73ba4..0a21a69 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,8 +1,10 @@ # Tesseract-JAX -```{include} content/get-started.md -:start-line: 2 -``` +`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. + +Now, learn how to [get started](content/get-started.md) with Tesseract-JAX, or explore the [API reference](content/api.md). ## License From b8bda0a9e869b235225ed9094217152e60de0a19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dion=20H=C3=A4fner?= Date: Wed, 23 Apr 2025 21:43:04 +0200 Subject: [PATCH 2/3] add git clone to instructions --- README.md | 3 ++- docs/content/get-started.md | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index b179849..363b64a 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,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: diff --git a/docs/content/get-started.md b/docs/content/get-started.md index 009d45c..6e60981 100644 --- a/docs/content/get-started.md +++ b/docs/content/get-started.md @@ -19,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: From 8b133657c2433e70d16010cb096bf1c16aa92d36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dion=20H=C3=A4fner?= Date: Wed, 23 Apr 2025 23:17:45 +0200 Subject: [PATCH 3/3] re-word introduction --- README.md | 13 +++++++++++-- docs/index.md | 15 ++++++++++++--- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 363b64a..ae5aa81 100644 --- a/README.md +++ b/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 diff --git a/docs/index.md b/docs/index.md index 0a21a69..f103d1d 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,10 +1,19 @@ # 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. +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. +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: -Now, learn how to [get started](content/get-started.md) with Tesseract-JAX, or explore the [API reference](content/api.md). +```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