bayesflow-org/bayesflow

Allow Pydantic models for network configs

Open

#344 opened on Mar 5, 2025

View on GitHub
 (1 comment) (0 reactions) (0 assignees)Python (85 forks)auto 404
featuregood first issuesugaruser interface

Repository metrics

Stars
 (694 stars)
PR merge metrics
 (PR metrics pending)

Description

Currently, the default configs are set like this

INTEGRATE_DEFAULT_CONFIG = {"method": "rk45", ..., tolerance=1e-3}

...

def __init__(self, ..., integrate_kwargs=None):
    self.integrate_kwargs = integrate_kwargs or INTEGRATE_DEFAULT_CONFIG

This presents several issues:

  1. There is no way to overwrite only part of the default config. Passing a dictionary for integrate_kwargs immediately overrides all other defaults.
  2. No validation is performed on the passed values, even if their values could only be relevant much later during training or at inference time, with no convenient way to fix them
  3. Some values, such as "method" are opaque with respect to what options are valid, and users need to rummage around in other parts of the docs

To fix, we would need to:

  1. Partially update the default configs with values the user passes
  2. Perform validation on the type and value of each parameter passed
  3. Improve our error messages, or provide a lot of redundant documentation

We could fix 1. and 2. with manual dictionary updates and parameter validation. However, using Pydantic's BaseModel class already allows both out of the box, while also fixing 3 with good error messages and highlighting from static type checkers.

This would change the above code to this:

class IntegrateConfig(BaseModel):
    method: Literal["rk45", "euler"] = "rk45"
    ...
    tolerance: PositiveFloat = 1e-3

...

def __init__(self, ..., integrate_kwargs=None):
    self.integrate_kwargs = FlowMatching.IntegrateConfig(**(integrate_kwargs or {}))

We could also change the name to (self.)integrate_config for consistency.

Disadvantages:

  • Increased maintenance for the configs
  • Slightly elevated entry barrier for contributors, since knowledge about Pydantic is required for correct typing

Contributor guide