{"id":19277527,"url":"https://github.com/enzymead/enzyme-jax","last_synced_at":"2026-02-17T00:03:02.350Z","repository":{"id":169909129,"uuid":"635473744","full_name":"EnzymeAD/Enzyme-JAX","owner":"EnzymeAD","description":"Custom Bindings for Enzyme Automatic Differentiation Tool and Interfacing with JAX.","archived":false,"fork":false,"pushed_at":"2025-04-05T19:44:49.000Z","size":6599,"stargazers_count":74,"open_issues_count":85,"forks_count":17,"subscribers_count":11,"default_branch":"main","last_synced_at":"2025-04-05T19:51:30.372Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"MLIR","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/EnzymeAD.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-05-02T19:16:08.000Z","updated_at":"2025-04-05T18:30:50.000Z","dependencies_parsed_at":null,"dependency_job_id":"2cda78e6-8067-4bee-9d65-3c4b176968a5","html_url":"https://github.com/EnzymeAD/Enzyme-JAX","commit_stats":null,"previous_names":["enzymead/enzyme-jax"],"tags_count":7,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/EnzymeAD%2FEnzyme-JAX","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/EnzymeAD%2FEnzyme-JAX/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/EnzymeAD%2FEnzyme-JAX/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/EnzymeAD%2FEnzyme-JAX/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/EnzymeAD","download_url":"https://codeload.github.com/EnzymeAD/Enzyme-JAX/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247987285,"owners_count":21028895,"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","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":[],"created_at":"2024-11-09T21:06:06.868Z","updated_at":"2026-02-17T00:03:02.334Z","avatar_url":"https://github.com/EnzymeAD.png","language":"MLIR","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Enzyme-JAX\n\nEnzyme-JAX is a C++ project whose original aim was to integrate the Enzyme automatic differentiation tool [1] with JAX, enabling automatic differentiation of external C++ code within JAX. It has since expanded to incorporate Polygeist's [2] high performance raising, parallelization, cross compilation workflow, as well as numerous tensor, linear algerba, and communication optimizations. The project uses LLVM's MLIR framework for intermediate representation and transformation of code. As Enzyme is language-agnostic, this can be extended for arbitrary programming\nlanguages (Julia, Swift, Fortran, Rust, and even Python)!\n\nYou can use\n\n```python\nfrom enzyme_ad.jax import cpp_call\n\n# Forward-mode C++ AD example\n\n@jax.jit\ndef something(inp):\n    y = cpp_call(inp, out_shapes=[jax.core.ShapedArray([2, 3], jnp.float32)], source=\"\"\"\n        template\u003cstd::size_t N, std::size_t M\u003e\n        void myfn(enzyme::tensor\u003cfloat, N, M\u003e\u0026 out0, const enzyme::tensor\u003cfloat, N, M\u003e\u0026 in0) {\n        out0 = 56.0f + in0(0, 0);\n        }\n        \"\"\", fn=\"myfn\")\n    return y\n\nones = jnp.ones((2, 3), jnp.float32)\nprimals, tangents = jax.jvp(something, (ones,), (ones,) )\n\n# Reverse-mode C++ AD example\n\nprimals, f_vjp = jax.vjp(something, ones)\n(grads,) = f_vjp((x,))\n```\n\n# Installation\n\nThe easiest way to install is using pip.\n\n```bash\n# The project is available on PyPi and installable like\n# a usual python package (https://pypi.org/project/enzyme-ad/)\npip install enzyme-ad\n```\n\n## Building from source\n\nRequirements: `bazel-6.5`, `clang++`, `python`, `python-virtualenv`,\n`python3-dev`.\n\nBuild our extension with:\n```sh\n# Will create a whl in bazel-bin/enzyme_ad-VERSION-SYSTEM.whl\nbazel build :wheel\n```\n\nFinally, install the built library with:\n```sh\npip install bazel-bin/enzyme_ad-VERSION-SYSTEM.whl\n```\nNote that you cannot run code from the root of the git directory. For instance, in the code below, you have to first run `cd test` before running `test.py`.\n\n## Running the test\n\nTo run tests, you can simply execute the following bazel commands (this does not require building or installing the wheel).\n```sh\nbazel test //test/...\n```\n\nAlternatively, if you have installed the wheel, you can manually invoke the tests as follows\n```sh\ncd test \u0026\u0026 python test.py\n```\n## LSP Support\n\nEnzyme-Jax exposes a bunch of different tensor rewrites as MLIR passes in `src/enzyme_ad/jax/Passes`. If you want to enable LSP support when working with this code, we recommend that you generate a `compile_commands.json` by running\n\n```bash\nbazel run :refresh_compile_commands\n```\n\n# References\n[1] Moses, William, and Valentin Churavy. \"Instead of rewriting foreign code for machine learning, automatically synthesize fast gradients.\" Advances in neural information processing systems 33 (2020): 12472-12485.\n\n[2] Moses, William S., et al. \"Polygeist: Raising C to polyhedral MLIR.\" 2021 30th International Conference on Parallel Architectures and Compilation Techniques (PACT). IEEE, 2021.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fenzymead%2Fenzyme-jax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fenzymead%2Fenzyme-jax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fenzymead%2Fenzyme-jax/lists"}