This is a library built on JAX and Flax NNX, implemented to create a new general framework for hypernetworks. Instead of training one neural network to approximate a single function, we train a network that generates the weights of another network on the fly, conditioned on a set of parameters, called hypervariables. This is particularly powerful for scientific computing, where one is often interested not in a single solution, but in how a solution evolves as physical parameters change.
The library is validated on two problems of increasing complexity:
toy— regression of analytic, parameter-dependent functions. Computationally soft to run on a laptop CPU, ideal to understand the method.turbulence— reconstruction of 2D fields from the turbulent radiative layer dataset (astrophysical turbulence), conditioned on time. A realistic, GPU-friendly workload.
Both experiments run with the same underlying code. This flexibility is mainly provided by the fact that everything (model topology, data source, optimizer, losses, metrics, post-processing) is declared in Hydra configuration files, so experiments are reproducible and can be reshaped from the command line without editing the source code.
IMPORTANT
To replicate code and run tests, follow instructions from Section §4, §5 and §6
- General Idea
- Repository structure
- Hydra
- Setup
- How to run tests
- General usage
- Inspecting Results
- Datasets
- Weights & Biases
- Running on an HPC cluster
A hypernetwork is a neural network whose output is the set of weights of a second network, the target network. In this library the data flow is always the same:
- Hypervariables (
z) are the parameters given as inputs to the hypernetwork. - Variables (
x) are the parameters given as inputs to the target network. - General pipeline:
- The Hypernetwork maps
zto a latent code - The projection heads expand that code into the actual kernels and biases of the target network
- The target network consumes
xwith those generated weights and produces the prediction.
- The Hypernetwork maps
The whole assembly is described as a small directed acyclic graph (DAG) of blocks and orchestrated
by the HypernetworkManager, which infers the execution order automatically. We provide two different families of neural netwoek architectures: plain MLPs and a SIRENs (sinusoidal representation network, useful when dealing with high-frequency signals).
| Folder | Purpose |
|---|---|
models/ |
The neural building blocks. Defines Hypernetwork, ProjectionHead/SirenHead (the heads that turn latent codes into weights), TargetNetwork (the network whose weights are injected rather than trained directly), and the HypernetworkManager that connect together all the blocks. Target architectures MLP and Siren are also implemented here. |
training/ |
The training engine, structured hierarchically like a Matryoshka doll. train_model drives the full execution, nesting train_epoch for individual epochs, which further nests train_step for single-batch execution (which is JIT-compiled). hypernet_utils.py manages the weight-injection mechanism (build_state_from_parameters, assign_parameters, apply). |
inference/ |
test_model — a single evaluation pass over a test set returning the required metrics. |
losses/ |
Composable loss functions: L2Loss (MSE), LpLoss (generic L-p norm) and CombinedLoss (a weighted sum of any of the above). |
metrics/ |
Flax-based streaming metrics: MSE, RMSE, RRMSE (relative RMSE) and MAE. For our results, we use RRMSE as most valuable performance metric. |
data_processing/ |
Dataset loading and construction. ToyDataSource generates samples by evaluating a parametric function over its domains; InMemoryHDF5Source loads a preprocessed HDF5 dataset into RAM. build_dataset wraps them in a Grain pipeline (shuffle/batch/iterate). For the turbulence problem, download_data.py fetches the dataset from Hugging Face, while preprocessing.py performs the temporal train/val/test split and normalization. |
postprocessing/ |
Plotting utilities: loss curves (plot_loss_curves), 1D/2D prediction plots, and plot_2d_hdf5_comparison (prediction vs. ground truth vs. error for the turbulence fields). |
utils/ |
General helpers. Most importantly the custom Hydra resolvers (hydra_resolvers.py), the safe math-expression parser get_function_from_string (used to define toy functions from strings), Orbax-based save_model/load_model, and learning-rate introspection. |
config/ |
The full Hydra configuration tree (see §3). This is where experiment parameters and configurations are defined. |
scripts/ |
Thin shell wrappers invoked by the Makefile. They allow the code to run with venv environment, Docker Images or on a OpenPBS based cluster. |
checkpoints/ |
Pre-trained model weights (downloaded from the GitHub release) used by the demonstration tests. |
results/ |
Default output directory. Each run writes here its logs, run_data.json, figures and a snapshot of the resolved config under .hydra/. |
main.py |
The single entry point. Reads the composed config, instantiates everything and runs the requested phases (train / test / inference plots). |
Hydra is the tool allowing flexibility in this project. The entire experiment setup is assembled from small, composable YAML files in config/ and any value can be overridden from the command line. Nothing in main.py is hard-coded.
config/
├── config.yaml # root: defaults, seed, flags (train/test/plot), W&B settings
├── problem/ # toy or turbulence problem
├── model/ # architecture choice
├── data_source/ # dataset definitions
├── toy_function/ # allows to construct different parametric functions
├── training/ # epochs, batch size, metrics, early stopping
├── optimizer/ # adam/adamw + schedulers
├── loss/ # l2.yaml | combined.yaml
├── preprocessing/ # turbulence dataset split & normalization
└── postprocessing/ # which plots to generate
The problem group is the controller: selecting problem=toy or problem=turbulence pulls in a set of data_source, training, model and postprocessing defaults consistent with the selected problem, so that the two workflows never get mixed.
To use our library and reproduce the results we provide two alternatives:
- a local
uvenvironment - a pre-built Docker image
Both are driven by the same Makefile, the only difference is the value of the ENGINE variable (venv or docker) and the fact that:
- Docker Image: Recommended for immediate reproduction. It arrives fully self-contained with all required datasets and model checkpoints pre-installed. No external setup or downloads are required.
- Local Setup: Requires you to manually load and place the datasets and checkpoints before running execution commands, as the local clone only contains the source code. This happens because hosting large binaries in Git violates repository best practices
NOTE: we use
uv, but a plainpython -m venvworks just as well.
-
Clone the repository:
git clone git@github.com:Riccardo-Rota/JAX_Hypernetworks.git
-
Move inside the repository root folder:
cd JAX_Hypernetworks -
Create the virtual environment with
uv(Python 3.12):uv venv jax_hypernetworks --python 3.12.3
-
Activate it:
source jax_hypernetworks/bin/activate -
Install the packages:
uv pip install -r requirements.txt
-
(Turbulence only) Load the astrophysical turbulence dataset from Hugging Face and run the provided preprocessing scripts:
make load-data
The toy problem needs no download: its data is generated on the fly.
-
(Optional) Download the pre-trained checkpoints used by the demonstration tests, from our
v1.0GitHub release:wget https://github.com/Riccardo-Rota/JAX_Hypernetworks/releases/download/v1.0/checkpoints.zip unzip checkpoints.zip rm checkpoints.zip
NOTE: this step is mandatory if you want to run provided tests.
-
Set the
ENGINEenvironment variable tovenvonce for all:export ENGINE="venv"
-
Run the code. You can now run experiments and the provided tests through the
Makefile. To run the provided tests, see section §5, while for general usage go to §6.
We also provide a ready-to-use Docker container, so you do not need to install any Python dependency yourself.
-
Clone the repository:
git clone git@github.com:Riccardo-Rota/JAX_Hypernetworks.git
-
Move inside the repository root folder:
cd JAX_Hypernetworks -
Pull the Docker image (install Docker and
docker loginfirst if needed):docker pull leonardobocchieri/jax-hypernetworks:latest
NOTE: the image is approximately 10 GB, so this step can take a few minutes.
-
Set the
ENGINEenvironment variable todocker(once and for all):export ENGINE="docker"
-
Run the code. You can now run experiments and the provided tests through the
Makefile. To run the provided tests, see section §5, while for general usage go to §6.
How the Docker engine works: what is baked in vs. mounted. The image contains the library code that needs to be kept unchanged (together with all dependencies). It is important to notice that
run_experiment.shbind-mounts the parts you actually edit between experiments:config/,main.pyandresults/. They are therefore taken from your local clone into the container. TheMakefileand thescripts/themselves run on the host and merely orchestrate thedocker runcall. The practical consequence: modifying a configuration file ormain.pylocally takes effect immediately inside the container. However, modifying local library files (such as themodels/directory) will have no effect on execution, because the container strictly executes the library copies baked into the Docker image.
Ownership note. Files created from inside the container belong to
root(specifically, theresults/folder). If you later run locally and hit a permission error, reclaim ownership withsudo chown -R $(whoami) results.
The execution device is selected with the USE_GPU flag of the Makefile (true / false, default
false). JAX is installed with CUDA 12 support.
- If you set
USE_GPU=truebut no NVIDIA driver/hardware is detected,run_experiment.shprints a hardware warning and safely reverts to CPU. - Independently,
main.pysetsJAX_PLATFORMS=cpuwhenever no GPU device is visible, so JAX itself never errors out.
The toy problem is small and is best run on CPU, while the turbulence problem benefits from a GPU.
Before proceeding, follow instructions in section §4
The Makefile provides a sequence of experimental tests that show the library capabilities. Each writes its output under a dedicated results/... directory. To understand how to read results, read Section §7
| Target | What it demonstrates |
|---|---|
make run-test1 |
(TOY) Inference from a good checkpoint. Loads a pre-trained MLP (train_model=false) and produces prediction and metric reports. The model should fit the target function well. |
make run-test2 |
(TOY) Training from scratch: the SIREN initialization matters. Trains a naive SIREN (poor initialization → poor fit), then a properly-initialized SIREN. This illustrates why SIREN's weight scheme is essential for high-frequency targets. |
make run-test3 |
(TOY) Fine-tuning a checkpoint. First loads a deliberately under-trained checkpoint (poor results), then resumes training (train_model=true) from those same weights to show the model recovering. |
make run-test4 |
(TOY) A challenging target. Trains a SIREN on a high-frequency sine wave (toy_function=highfreq_sine). This illustrates how SIREN architectures perform well on high-frequency signals. |
make run-test5 |
(PERFORMANCE) JIT vs. no-JIT. Runs the same short training twice on CPU (with and without JAX JIT compilation) and prints the measured speed-up. |
make run-test6 |
(TURBULENCE) Inference with MLPs. Loads a pre-trained turbulence MLP and reconstructs the velocity fields, comparing prediction against ground truth. Requires the turbulence dataset (make load-data). |
make run-test7 |
(TURBULENCE) Inference with SIRENs. Same as test 6 but with a SIREN target on the density field. Requires the turbulence dataset (make load-data). |
The library allows to run personalized experiments. Every experiment is launched through make run-local, which accepts three knobs:
make run-local USE_GPU=<true|false> ENGINE=<venv|docker> OVERRIDES="<hydra overrides>"ENGINE—venv(local interpreter) ordocker.USE_GPU—trueorfalse(defaultfalse). See the CPU/GPU flag note above.OVERRIDES— any space-separated list of Hydra overrides. Alternatively, edit the files inconfig/directly.
Full example:
make run-local USE_GPU=true OVERRIDES="problem=toy model=toy_siren training.epochs=1000"main.py executes up to three phases, each controlled by a boolean flag (overridable):
train_model=true: train and checkpoint the model.test_model=true: evaluate metrics on the test set.plot_inference=true: generate the post-processing figures.
To merely inspect a checkpoint, set train_model=false checkpoint='checkpoints/<name>', exactly as the
demonstration tests do.
When running on GPU you may see repeated lines such as:
W external/xla/.../dot_search_space.cc:200] All configs were filtered out because none of them sufficiently match the hints. ... Working around this by using the full hints set instead.These are warnings, not errors — the run is fine. The warning concerns only how XLA chose the kernel and it happens once, during compilation. It depends on your GPU and JAX version, so it may not show up on every machine. To silence them, set
TF_CPP_MIN_LOG_LEVEL=2before launching.
Results are saved in the results/ folder, which is structured as follows:
| Path | What it contains |
|---|---|
.hydra/ |
The YAML files with all the configuration settings used for that specific run. |
.checkpoints/ |
Model checkpoints folders. |
.figures/ |
Plots and Loss curves in PNG format. |
default.log |
CLI output prints. |
run:data.json |
General results for metrics, earely stopping, best epoch and training history. |
training_log.txt |
Training log with metrics and results after each epoch. |
No download is required. ToyDataSource generates samples on demand by evaluating a parametric function
(parsed safely from a string via get_function_from_string) over its hyper- and variable domains. The
function and domains are defined in config/toy_function/ — default.yaml for the standard target and
highfreq_sine.yaml for the hard, high-frequency case.
The turbulence workload uses the turbulent radiative layer 2D dataset
(polymathic-ai/turbulent_radiative_layer_2D
on Hugging Face). make load-data runs two steps:
download_data.py— fetches the raw HDF5 file for the selected cooling timetcool(configurable inconfig/preprocessing/turbulence.yaml).preprocessing.py— performs a temporal train/val/test split (consecutive-timestep chunks are allocated to each split), computes normalization statistics from the training split only, and writesturbulence_dataset_{train,val,test}.hdf5intodatasets/.
At run time InMemoryHDF5Source loads the relevant split fully into RAM. The fields, coordinates and
prediction targets (density, pressure, velocity_x, velocity_y) are declared in
config/data_source/turbulence.yaml and can be remapped from the command line.
Integrating a custom dataset requires no changes to the underlying Python codebase. You only need to provide pre-split .hdf5 files and map their structure using a YAML configuration.
1. Data Placement
Place your .hdf5 files directly into the datasets/ folder. Your data must already be divided into training, validation and test splits (e.g., custom_train.hdf5, custom_val.hdf5, custom_test.hdf5).
2. Configuration Mapping
Create a new YAML file inside config/data_source/ (e.g., custom_data.yaml). This file dictates how the data loader parses your .hdf5 columns. You must define a schema that maps your dataset's variable names to their exact column indices. Using those names, you then categorize the columns into hypervariables (hypervar_keys), physical coordinates/inputs (var_keys), and prediction targets (target_keys). Check config/data_source/turbulence.yaml to have a deeper insight.
The project integrates with Weights & Biases for live monitoring and logging. It is disabled by default, but it can be turned on by adding use_wandb=true to overrides. Make sure wandb_settings are set, either editing config/config.yaml file or with:
make run-local OVERRIDES="problem=toy use_wandb=true wandb_settings.project=my_project wandb_settings.entity=my_team"Generate an API key from your W&B account and store it securely in a ~/.netrc file in your home
directory:
touch ~/.netrc
chmod 600 ~/.netrcThen open the file in a text editor (e.g. nano ~/.netrc) and paste the following block:
machine api.wandb.ai
login user
password <your_API_key>
(Replace <your_API_key> with your actual 40-character Weights & Biases token. Do not change the word
user.)
run_experiment.sh parses this file automatically: if a key is found it logs you in; if not, it forces offline mode so a run never blocks on a login prompt.
Runs produced offline (typically on a cluster) can be uploaded afterwards:
RESULTS_DIR=results/runs_turbulence make sync-clusterThis section concerns Politecnico di Milano cluster. If you do not have access to it, you can skip the reading.
For long or GPU-heavy jobs, the library runs on an HPC cluster via apptainer (an open-source container platform designed to run on HPC systems) and the PBS scheduler.
- Clone the repository on the cluster (preferably under your
workspace). - Convert the Docker image to an apptainer
.sif:apptainer pull jax-hypernetworks.sif docker://leonardobocchieri/jax-hypernetworks:latest
- Submit a job with
make submit-cluster. It accepts two extra arguments:USE_GPU(true/false) — decides which queue the job is sent to (GPU or CPU).OVERRIDES— the Hydra overrides string (or edit the files inconfig/directly).
make submit-cluster USE_GPU=true OVERRIDES="problem=turbulence use_wandb=true training.epochs=500"
For other relevant information, consult the student cluster guidelines on Webeep.
See LICENSE.