Skip to content

Installation

This page covers the supported ways to install jNO and set up an environment for running examples and training models.

Prerequisites

  • Python 3.11 - 3.13
  • uv (recommended package manager) or conda or pip (alternative package managers)

Conda Environment

Create and activate a fresh conda environment first if you want an isolated Python setup:

conda create -n jno python=3.12
conda activate jno

pip Installation

Install jNO directly with pip:

pip install jax-neural-operators

If a Nvidia GPU is avaiable install:

pip install jax[cuda]

For editable local development from a cloned repository:

pip install -e .

Additionally, install the dev extra for development dependencies:

pip install jax-neural-operators[dev]

For deployment scenarios, jNO supports exporting compiled models via IREE: 

pip install jax-neural-operators[iree]

Installation with uv

jNO uses uv for environment and dependency management. uv installs packages into a virtual environment under your user directory, so no sudo is required.

CPU (default)

uv sync

GPU / CUDA

uv sync --extra cuda

Developer tools

uv sync --extra dev

IREE ahead-of-time compiler

uv sync --extra iree

Windows - allow script execution first

Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser

Next Step

After installation, continue with Getting Started to run your first example and learn the basic jNO workflow.