eqxconfig

eqxconfig is a small utility for constructing an Equinox module from TOML without writing a second configuration class. It derives a Pydantic model from the concrete module's annotated __init__ signature, validates the configuration, and calls the unchanged constructor.

TOML
  ↓
dynamically generated Pydantic schema
  ↓
validated constructor arguments
  ↓
Equinox __init__
  ↓
ordinary eqx.Module / JAX PyTree

The package requires Python 3.14 or newer.

Installation

python -m pip install eqxconfig

Usage

import equinox as eqx
import jax

from eqxconfig import ConfigurableModule


class MLP(ConfigurableModule):
    mlp: eqx.nn.MLP

    def __init__(
        self,
        in_size: int,
        out_size: int,
        width_size: int = 128,
        depth: int = 3,
        seed: int = 0,
    ) -> None:
        self.mlp = eqx.nn.MLP(
            in_size=in_size,
            out_size=out_size,
            width_size=width_size,
            depth=depth,
            key=jax.random.key(seed),
        )


model = MLP.from_toml("model.toml")
model_with_overrides = MLP.from_toml(
    "model.toml",
    width_size=512,
    depth=6,
)

model.toml contains constructor arguments at its root:

in_size = 32
out_size = 10
width_size = 256
depth = 4
seed = 42

Precedence is:

__init__ defaults < TOML < explicit from_toml overrides

The final configuration—including constructor defaults, TOML values, and explicit overrides—goes through one Pydantic v2 validation pass. The generated schema preserves annotations and defaults, so Annotated/Field constraints, Literal, unions, enums, constrained types, and nested Pydantic models work normally. Unknown keys are forbidden. Pydantic validation errors, TOML decoding errors, and file access errors retain their standard exception types.

Nested ConfigurableModule fields are derived recursively from their own constructor signatures. Each nested value is built through its ordinary Equinox constructor, so custom __init__ methods, field converters, __post_init__, __check_init__, frozen-module semantics, and PyTree behavior are preserved at every level. This also applies inside optional, list, tuple, and dictionary fields. Passing an already constructed nested module as an override preserves that instance without constructing or checking it again.

The schema is generated lazily once per Equinox subclass. Constructor parameters must be annotated and must be positional-or-keyword or keyword-only. Positional- only parameters, *args, and **kwargs are deliberately rejected with a TypeError. Constructor defaults must be ordinary Python defaults in the signature; defaults embedded in Field(...), private parameter names, and the reserved name model_config are rejected because they cannot preserve a one-to-one constructor schema. Configuration keys match constructor parameter names; Pydantic validation aliases are not supported in v0.1.

Direct construction remains ordinary Python and Equinox:

model = MLP(in_size=32, out_size=10, width_size=512)

Configuration parsing and module construction happen before JAX transformations. Do not call from_toml inside jax.jit; pass the resulting module into compiled functions as a normal Equinox PyTree.

Development

Install the development dependency group with uv:

uv sync --all-groups

Run the individual checks with:

uv run ruff check .
uv run ruff format --check .
uv run pyrefly check
uv run pytest
uv build
uv run twine check dist/*

The repository uses prek for local hooks. From a fresh checkout:

prek install
prek run --all-files

Releases

The package version lives in pyproject.toml. Update it with uv version, merge the change, and push a matching tag such as v0.1.0. The release workflow builds and validates one wheel and one source distribution, attaches those exact files to a GitHub Release, and publishes the same files to PyPI.

PyPI publishing uses Trusted Publishing, not an API token. Before the first release, the repository owner must:

1. Create a GitHub environment named pypi (and add any desired deployment protection rules).
2. In the PyPI project's publishing settings, add a GitHub Trusted Publisher for the actual repository owner and repository name, workflow file release.yml, and environment pypi.

No long-lived PyPI credential should be added to GitHub secrets.

Scope

eqxconfig intentionally handles only root-level TOML-to-constructor configuration. It is not a general configuration framework and does not add environment variables, command-line parsing, remote sources, or configuration work inside JAX transformations.