https://github.com/hayden-donnelly/jax-to-cpp

Generate HLO with JAX then execute in C++.
https://github.com/hayden-donnelly/jax-to-cpp

Last synced: about 10 hours ago
JSON representation

Generate HLO with JAX then execute in C++.

• Host: GitHub
• URL: https://github.com/hayden-donnelly/jax-to-cpp
• Owner: hayden-donnelly
• License: apache-2.0
• Created: 2023-11-15T23:17:28.000Z (about 1 year ago)
• Default Branch: main
• Last Pushed: 2024-10-26T19:39:09.000Z (24 days ago)
• Last Synced: 2024-10-26T21:45:26.623Z (24 days ago)
• Language: C++
• Size: 133 KB
• Stars: 2
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE.txt

Awesome Lists containing this project

README

        






# jax-to-cpp

A demo showing how you can generate HLO with JAX then execute in C++. This removes the overhead of Python at runtime while maintaining its flexibility at design-time.

## Getting Started

Recursively clone this repository to get all the necessary files (including the [XLA submodule](https://github.com/openxla/xla/tree/dd7b496b6d082bfeda230332ea8e418e8cf3ca78)):

```

git clone --recursive https://github.com/hayden-donnelly/jax-to-cpp.git

```

Change into the cloned directory:

```

cd jax-to-cpp

```

Build the Docker image:

```

docker-compose build

```

Start the Docker container:

```

docker-compose up -d

```

Open a shell in the container:

```

docker-compose exec j2cpp bash

```

Generate a simple HLO module with JAX:

```

python3 scripts/generate_hlo.py

```

Build the ``execute_hlo`` target with Bazel (this might take a while at first, but the results will be cached for next time):

```

bazel build //:execute_hlo

```

Run the executable you just built:

```

bazel-bin/execute_hlo

```