Skip to content

Commit 77d4f7f

Browse files
dionhaefnerdionhaefner
committed
add missing files
1 parent c8ba7a2 commit 77d4f7f

File tree

3 files changed

+88
-30
lines changed

3 files changed

+88
-30
lines changed

‎demo/cfd/demo.ipynb‎

Lines changed: 4 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
"output_type": "stream",
2020
"text": [
2121
"\u001b[2K \u001b[1;2m[\u001b[0m\u001b[34mi\u001b[0m\u001b[1;2m]\u001b[0m Building image \u001b[33m...\u001b[0m\n",
22-
"\u001b[2K\u001b[37m⠇\u001b[0m \u001b[37mProcessing\u001b[0m\n",
23-
"\u001b[1A\u001b[2K \u001b[1;2m[\u001b[0m\u001b[34mi\u001b[0m\u001b[1;2m]\u001b[0m Built image sh\u001b[1;92ma256:b5e2\u001b[0m455fe251, \u001b[1m[\u001b[0m\u001b[32m'jax-cfd:latest'\u001b[0m\u001b[1m]\u001b[0m\n"
22+
"\u001b[2K\u001b[37m⠋\u001b[0m \u001b[37mProcessing\u001b[0m\n",
23+
"\u001b[1A\u001b[2K \u001b[1;2m[\u001b[0m\u001b[34mi\u001b[0m\u001b[1;2m]\u001b[0m Built image sh\u001b[1;92ma256:0c56\u001b[0m253398c1, \u001b[1m[\u001b[0m\u001b[32m'jax-cfd:latest'\u001b[0m\u001b[1m]\u001b[0m\n"
2424
]
2525
},
2626
{
@@ -46,32 +46,6 @@
4646
"name": "stdout",
4747
"output_type": "stream",
4848
"text": [
49-
"Looking in indexes: https://pkgs.dev.azure.com/pasteur-labs/d89796ea-4a5b-48aa-9930-4cebcdc9d64a/_packaging/internal/pypi/simple, https://pypi.org/simple\n",
50-
"Requirement already satisfied: jax_cfd in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from -r requirements.txt (line 1)) (0.2.1)\n",
51-
"Requirement already satisfied: matplotlib in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from -r requirements.txt (line 2)) (3.10.0)\n",
52-
"Requirement already satisfied: pillow in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from -r requirements.txt (line 3)) (11.1.0)\n",
53-
"Requirement already satisfied: optax in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from -r requirements.txt (line 4)) (0.2.4)\n",
54-
"Requirement already satisfied: jax in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from jax_cfd->-r requirements.txt (line 1)) (0.5.2)\n",
55-
"Requirement already satisfied: numpy in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from jax_cfd->-r requirements.txt (line 1)) (2.2.0)\n",
56-
"Requirement already satisfied: scipy in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from jax_cfd->-r requirements.txt (line 1)) (1.15.0)\n",
57-
"Requirement already satisfied: tree-math in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from jax_cfd->-r requirements.txt (line 1)) (0.2.1)\n",
58-
"Requirement already satisfied: contourpy>=1.0.1 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from matplotlib->-r requirements.txt (line 2)) (1.3.1)\n",
59-
"Requirement already satisfied: cycler>=0.10 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from matplotlib->-r requirements.txt (line 2)) (0.12.1)\n",
60-
"Requirement already satisfied: fonttools>=4.22.0 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from matplotlib->-r requirements.txt (line 2)) (4.55.8)\n",
61-
"Requirement already satisfied: kiwisolver>=1.3.1 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from matplotlib->-r requirements.txt (line 2)) (1.4.8)\n",
62-
"Requirement already satisfied: packaging>=20.0 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from matplotlib->-r requirements.txt (line 2)) (24.2)\n",
63-
"Requirement already satisfied: pyparsing>=2.3.1 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from matplotlib->-r requirements.txt (line 2)) (3.2.1)\n",
64-
"Requirement already satisfied: python-dateutil>=2.7 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from matplotlib->-r requirements.txt (line 2)) (2.9.0.post0)\n",
65-
"Requirement already satisfied: absl-py>=0.7.1 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from optax->-r requirements.txt (line 4)) (2.1.0)\n",
66-
"Requirement already satisfied: chex>=0.1.87 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from optax->-r requirements.txt (line 4)) (0.1.88)\n",
67-
"Requirement already satisfied: jaxlib>=0.4.27 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from optax->-r requirements.txt (line 4)) (0.5.1)\n",
68-
"Requirement already satisfied: etils[epy] in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from optax->-r requirements.txt (line 4)) (1.11.0)\n",
69-
"Requirement already satisfied: typing_extensions>=4.2.0 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from chex>=0.1.87->optax->-r requirements.txt (line 4)) (4.12.2)\n",
70-
"Requirement already satisfied: setuptools in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from chex>=0.1.87->optax->-r requirements.txt (line 4)) (75.6.0)\n",
71-
"Requirement already satisfied: toolz>=0.9.0 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from chex>=0.1.87->optax->-r requirements.txt (line 4)) (1.0.0)\n",
72-
"Requirement already satisfied: ml_dtypes>=0.4.0 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from jax->jax_cfd->-r requirements.txt (line 1)) (0.5.1)\n",
73-
"Requirement already satisfied: opt_einsum in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from jax->jax_cfd->-r requirements.txt (line 1)) (3.4.0)\n",
74-
"Requirement already satisfied: six>=1.5 in /Users/dion/.virtualenvs/science/lib/python3.12/site-packages (from python-dateutil>=2.7->matplotlib->-r requirements.txt (line 2)) (1.17.0)\n",
7549
"\n",
7650
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\n",
7751
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
@@ -81,7 +55,7 @@
8155
],
8256
"source": [
8357
"# Install additional requirements\n",
84-
"%pip install -r requirements.txt"
58+
"%pip install -r requirements.txt -q"
8559
]
8660
},
8761
{
@@ -852,7 +826,7 @@
852826
},
853827
{
854828
"cell_type": "code",
855-
"execution_count": null,
829+
"execution_count": 12,
856830
"metadata": {},
857831
"outputs": [],
858832
"source": [

‎docs/content/index.md‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Tesseract-JAX
2+
3+
`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.
4+
5+
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.
6+
7+
## Quick start
8+
9+
```{note}
10+
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+).
11+
```
12+
13+
```{seealso}
14+
For more detailed installation instructions, please refer to the [Tesseract Core documentation](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/introduction/installation.html).
15+
```
16+
17+
1. Install Tesseract-JAX:
18+
19+
```bash
20+
$ pip install tesseract-jax
21+
```
22+
23+
2. Build an example Tesseract:
24+
25+
```bash
26+
$ tesseract build demo/simple/vectoradd_jax
27+
```
28+
29+
3. Use it as part of a JAX program:
30+
31+
```python
32+
import jax
33+
import jax.numpy as jnp
34+
from tesseract_core import Tesseract
35+
from tesseract_jax import apply_tesseract
36+
37+
# Load the Tesseract
38+
t = Tesseract.from_image("vectoradd_jax")
39+
40+
# Run it with JAX
41+
x = jnp.ones((1000, 1000))
42+
y = jnp.ones((1000, 1000))
43+
44+
def vector_add(x, y):
45+
return apply_tesseract(t, x, y)
46+
47+
vector_add(x, y) # success!
48+
49+
# You can also use it with JAX transformations like JIT and grad
50+
vector_add_jit = jax.jit(vector_add)
51+
vector_add_jit(x, y)
52+
53+
vector_add_grad = jax.grad(vector_add)
54+
vector_add_grad(x, y)
55+
```
56+
57+
```{tip}
58+
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.
59+
```
60+
61+
## Sharp edges
62+
63+
- **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.
64+
65+
```python
66+
from tesseract_core import Tesseract
67+
from tesseract_jax import apply_tesseract
68+
69+
tess = Tesseract.from_image("vectoradd")
70+
apply_tesseract(tess, {"a": 1.0, "b": 2.0}) # ❌ raises an error
71+
apply_tesseract(tess, {"a": jnp.array(1.0), "b": jnp.array(2.0)}) # ✅ works
72+
```
73+
- **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.
74+
75+
```{tip}
76+
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`.
77+
```
78+
79+
## License
80+
81+
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).
82+
83+
Tesseract is a registered trademark of Pasteur Labs, Inc. and may not be used without permission.

‎docs/demo_notebooks/.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
*.ipynb

0 commit comments

Comments
 (0)