{"id":30897600,"url":"https://github.com/pasteurlabs/tesseract-jax","last_synced_at":"2026-04-03T21:11:56.430Z","repository":{"id":289833460,"uuid":"969975613","full_name":"pasteurlabs/tesseract-jax","owner":"pasteurlabs","description":"Execute + differentiate Tesseracts as part of JAX programs, with full support for function transformations like JIT, grad, and more. ⚡ ","archived":false,"fork":false,"pushed_at":"2025-09-04T11:26:41.000Z","size":11598,"stargazers_count":21,"open_issues_count":8,"forks_count":1,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-09-05T01:52:49.509Z","etag":null,"topics":["automatic-differentiation","differentiable-programming","jax","remote-procedure-call","scientific-machine-learning","tesseract"],"latest_commit_sha":null,"homepage":"https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/pasteurlabs.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":null,"codeowners":".github/CODEOWNERS","security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2025-04-21T08:46:49.000Z","updated_at":"2025-09-04T11:26:29.000Z","dependencies_parsed_at":null,"dependency_job_id":"66defdfb-91c0-4443-a7e3-192c10bd6a2c","html_url":"https://github.com/pasteurlabs/tesseract-jax","commit_stats":null,"previous_names":["pasteurlabs/tesseract-jax"],"tags_count":2,"template":false,"template_full_name":null,"purl":"pkg:github/pasteurlabs/tesseract-jax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pasteurlabs%2Ftesseract-jax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pasteurlabs%2Ftesseract-jax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pasteurlabs%2Ftesseract-jax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pasteurlabs%2Ftesseract-jax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/pasteurlabs","download_url":"https://codeload.github.com/pasteurlabs/tesseract-jax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/pasteurlabs%2Ftesseract-jax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":274231110,"owners_count":25245685,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-09-08T02:00:09.813Z","response_time":121,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["automatic-differentiation","differentiable-programming","jax","remote-procedure-call","scientific-machine-learning","tesseract"],"created_at":"2025-09-09T00:13:51.483Z","updated_at":"2026-04-03T21:11:56.425Z","avatar_url":"https://github.com/pasteurlabs.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"docs/static/logo-transparent.png\" width=\"128\" align=\"right\"\u003e\n\n### Tesseract-JAX\n\nTesseract-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.\n\n[Read the docs](https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/) |\n[Explore the examples](https://github.com/pasteurlabs/tesseract-jax/tree/main/examples) |\n[Report an issue](https://github.com/pasteurlabs/tesseract-jax/issues) |\n[Talk to the community](https://si-tesseract.discourse.group/) |\n[Contribute](CONTRIBUTING.md)\n\n---\n\nThe 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:\n\n```python\n@jax.jit\ndef vector_sum(x, y):\n    res = apply_tesseract(vectoradd_tesseract, {\"a\": {\"v\": x}, \"b\": {\"v\": y}})\n    return res[\"vector_add\"][\"result\"].sum()\n\njax.grad(vector_sum)(x, y) # 🎉\n```\n\n## Quick start\n\n\u003e [!NOTE]\n\u003e 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+).\n\n\u003e [!IMPORTANT]\n\u003e For more detailed installation instructions, please refer to the [Tesseract Core documentation](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/introduction/installation.html).\n\n1. Install Tesseract-JAX:\n\n   ```bash\n   $ pip install tesseract-jax\n   ```\n\n2. Build an example Tesseract:\n\n   ```bash\n   $ git clone https://github.com/pasteurlabs/tesseract-jax\n   $ tesseract build tesseract-jax/examples/simple/vectoradd_jax\n   ```\n\n3. Use it as part of a JAX program via the JAX-native `apply_tesseract` function:\n\n   ```python\n   import jax\n   import jax.numpy as jnp\n   from tesseract_core import Tesseract\n   from tesseract_jax import apply_tesseract\n\n   # Load the Tesseract\n   t = Tesseract.from_image(\"vectoradd_jax\")\n   t.serve()\n\n   # Run it with JAX\n   x = jnp.ones((1000,))\n   y = jnp.ones((1000,))\n\n   def vector_sum(x, y):\n       res = apply_tesseract(t, {\"a\": {\"v\": x}, \"b\": {\"v\": y}}, vmap_method=\"sequential\")\n       return res[\"vector_add\"][\"result\"].sum()\n\n   vector_sum(x, y) # success!\n\n   # You can also use it with JAX transformations like JIT and grad\n   vector_sum_jit = jax.jit(vector_sum)\n   vector_sum_jit(x, y)\n\n   vector_sum_grad = jax.grad(vector_sum)\n   vector_sum_grad(x, y)\n\n   # vmap requires an explicit vmap_method — \"sequential\" is safe but slow\n   # while \"auto_experimental\" or \"expand_dims\" is more efficient for Tesseracts that support batching.\n   # See https://docs.pasteurlabs.ai/projects/tesseract-jax/latest/content/vmap-methods.html\n   vector_sum_vmap = jax.vmap(vector_sum)\n   vector_sum_vmap(x.reshape(10, 100), y.reshape(10, 100))\n   ```\n\n\u003e [!TIP]\n\u003e 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.\n\n## Sharp edges\n\n- **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 when used in conjunction with automatic differentiation and JAX transformations. This is because JAX, in these cases, 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.\n\n\u003e [!TIP]\n\u003e 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`.\n\n## License\n\nTesseract-JAX is licensed under the [Apache License 2.0](LICENSE) and is free to use, modify, and distribute (under the terms of the license).\n\nTesseract is a registered trademark of Pasteur Labs, Inc. and may not be used without permission.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpasteurlabs%2Ftesseract-jax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpasteurlabs%2Ftesseract-jax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpasteurlabs%2Ftesseract-jax/lists"}