annotated-jax

Deep Learning examples using the Jax ecosystem of libraries

env setup

Important

Ensure that you've uv installed: link.

uv sync
source .venv/bin/activate

Tip

You can also use the provided devcontainer configuration.

devcontainer up --workspace-folder .

run jupyter lab

# inside the project root
jupyter lab