API Reference

Hybrid Model

Ensemble module: HybridModel combining neural networks and physics solvers.

class adeptml.ensemble.HybridModel(*args: Any, **kwargs: Any)

Torch module for serial hybrid physics-informed models.

Combines neural network modules (MLPs, custom torch.nn.Module subclasses) with non-differentiable physics models in a single differentiable computation graph. Models are executed in the insertion order of config.models.

Parameters:

config (HybridConfig) – Configuration object specifying all sub-models and optional input routing. See HybridConfig for details.

Examples

Serial NN → physics pipeline:

cfg = HybridConfig(models={"nn": mlp_cfg, "physics": phy_cfg})
model = HybridModel(cfg)
out = model(x, phy_args=[arg1, arg2])

Custom input routing — pass original input to the physics model regardless of what the NN outputs:

cfg = HybridConfig(
    models={"nn": mlp_cfg, "physics": phy_cfg},
    model_inputs={"physics": {"Input": None, "nn": None}},
)

forward(x, phy_args=None)

Run inference on the hybrid model.

Parameters:

• x (torch.Tensor) – Input tensor (batch, in_dim).

• phy_args – Extra positional arguments forwarded to physics sub-models (e.g. [arg1, arg2]). May be a list of tensors or None.

Returns:

Output of the final model in the pipeline.

Return type:

torch.Tensor

Models

AdePT-ML provides three autograd-compatible wrappers for non-differentiable physics functions, differentiated by how gradients are computed during the backward pass:

Class	forward_func signature	jacobian_func signature
Physics	(x, *args) -> ndarray	(x, *args) -> ndarray shape (batch, out, in)
Physics_VJP	(x, *args) -> ndarray	(x, grad_out, *args) -> ndarray shape (batch, in)
Physics_SplitVJP	(x, *args) -> (ndarray, pullback_fn)	not used

When to use each:

• Use Physics when the full Jacobian is cheap or already available (e.g. small output dimension, analytical Jacobian).

• Use Physics_VJP when a manual VJP avoids materialising the full Jacobian (e.g. adjoint method, large output dim).

• Use Physics_SplitVJP when the physics solver already produces a pullback closure during the forward pass (e.g. via jax.vjp). This is the most memory-efficient option because no part of the physics forward is re-run during backpropagation.

  Example with a JAX solver:

  from my_physics import physics_forward
  # physics_forward(x, *args) -> (y, pullback_fn)
  
  phy_cfg = PhysicsConfig(
      forward_func=physics_forward,
      use_split_vjp=True,
  )
  

class adeptml.models.MLP(*args: Any, **kwargs: Any)

Multilayer Perceptron (MLP) neural network model.

config

Instance of MLPConfig.

Type:

MLPConfig

Note

This class implements a Multilayer Perceptron (MLP) neural network model. It takes a configuration dataclass with parameters such as hidden layer size, input and output dimensions, number of hidden layers, and activation functions.

forward(x)

Forward pass of the MLP model.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor

class adeptml.models.Physics(*args: Any, **kwargs: Any)

Custom autograd function wrapping a physics model with a full Jacobian.

The backward pass materialises the full (batch, out, in) Jacobian and contracts it with the upstream gradient via a batched matrix-vector product. Use Physics_VJP when a manual VJP is cheaper than the full Jacobian, or Physics_SplitVJP when the physics solver already produces a pullback closure (e.g. via jax.vjp).

See also

Physics_VJP, Physics_SplitVJP

static backward(ctx, grad_output)

Compute VJP via full Jacobian: grad = grad_output @ J.

static forward(ctx, x: torch.Tensor, forward_fun: Callable, jacobian_fun: Callable, args: List[torch.Tensor] | None = None)

Run the physics forward pass.

Parameters:

• ctx – PyTorch autograd context.

• x – Input tensor (batch, in_dim).

• forward_fun – Physics function (x_np, *args_np) -> ndarray.

• jacobian_fun – Jacobian function (x_np, *args_np) -> ndarray with shape (batch, out_dim, in_dim).

• args – Extra positional arguments passed to forward_fun and jacobian_fun. Gradients are not computed w.r.t. these.

Returns:

Output tensor (batch, out_dim).

Return type:

torch.Tensor

class adeptml.models.Physics_SplitVJP(*args: Any, **kwargs: Any)

Custom autograd function for physics models that expose a pullback closure.

This mode supports forward_func functions that return both the output and a pullback closure — the pattern produced by jax.vjp:

y, pullback_fn = forward_func(x, *args)
# later, during backward:
(dx, *_) = pullback_fn(grad_output)

The pullback closure is stored in the autograd context during the forward pass and invoked during the backward pass. No part of the physics forward is re-executed during backpropagation, making this the most memory-efficient mode when the solver is expensive and already caches its intermediates (e.g. via jax.checkpoint inside jax.vjp).

jacobian_func is not used in this mode and should be omitted from PhysicsConfig.

Example

Given the JAX physics module:

from Funcs import physics
# physics(x)-> (y_normalized, pullback_fn)

Configure and use as:

from adeptml.configs import PhysicsConfig, HybridConfig
from adeptml.ensemble import HybridModel

phy_cfg = PhysicsConfig(
    forward_func=physics,
    use_split_vjp=True,
)
hybrid_cfg = HybridConfig(models={"nn": mlp_cfg, "physics": phy_cfg})
model = HybridModel(hybrid_cfg)

# During training the pullback is cached automatically; no extra code
# is needed in the training loop.

See also

Physics, Physics_VJP

static backward(ctx, grad_output)

Call the stored pullback closure with the upstream gradient.

The cotangent for x is pullback_fn(grad_output)[0]; cotangents for forward_fun and args are None.

static forward(ctx, x: torch.Tensor, forward_fun: Callable, args: List[torch.Tensor] | None = None)

Run the split-VJP physics forward pass.

Calls forward_fun(x_np, *args_np) which must return (output_np, pullback_fn). The pullback is stored in ctx for use during backward().

Parameters:

• ctx – PyTorch autograd context.

• x – Input tensor (batch, in_dim).

• forward_fun – Split-VJP physics function (x_np, *args_np) -> (ndarray, callable).

• args – Extra positional arguments (e.g. boundary conditions, initial conditions, simulation time). Gradients are not computed w.r.t. these.

Returns:

Output tensor (batch, out_dim).

Return type:

torch.Tensor

class adeptml.models.Physics_VJP(*args: Any, **kwargs: Any)

Custom autograd function wrapping a physics model with a manual VJP.

Unlike Physics, the backward pass delegates to a user-supplied jacobian_func(x, grad_output, *args) -> ndarray that returns the VJP directly, avoiding explicit Jacobian materialisation.

See also

Physics, Physics_SplitVJP

static backward(ctx, grad_output)

Compute VJP by calling the user-supplied VJP function directly.

static forward(ctx, x: torch.Tensor, forward_fun: Callable, jacobian_fun: Callable, args: List[torch.Tensor] | None = None)

Run the physics forward pass.

Parameters:

• ctx – PyTorch autograd context.

• x – Input tensor (batch, in_dim).

• forward_fun – Physics function (x_np, *args_np) -> ndarray.

• jacobian_fun – VJP function (x_np, grad_output_np, *args_np) -> ndarray with shape (batch, in_dim).

• args – Extra positional arguments. Gradients are not computed w.r.t. these.

Returns:

Output tensor (batch, out_dim).

Return type:

torch.Tensor

Configs

class adeptml.configs.HybridConfig(models: dict[str, MLPConfig | PhysicsConfig | torch.nn.Module], model_inputs: dict[str, dict[str, list[int] | None]] | None = None)

Configuration for Hybrid (ensemble) models.

models

Maps model names (str) to a ModelConfig or an existing torch.nn.Module. Models are executed in insertion order.

Type:

dict

model_inputs

By default the ensemble runs sequentially — the output of each model becomes the input of the next. Setting this dict overrides that behaviour for specific models.

Keys are model names; values are dicts mapping source model names to dimension selectors:

• None — pass the entire source tensor.

• [int, ...] — pass only the listed dimensions (tensor[:, dims]).

Use the special key "Input" to reference the original input to HybridModel.forward().

Example — concatenate dims 0-3 of the original input with the full output of "mlp1":

model_inputs={
    "physics": {
        "Input": [0, 1, 2, 3],
        "mlp1": None,
    }
}

Type:

dict, optional

class adeptml.configs.MLPConfig(num_input_dim: int, num_hidden_dim: int, num_output_dim: int, num_hidden_layers: int, hidden_activation: str, output_activation: str)

Configuration class for the Multilayer Perceptron (MLP) model.

num_input_dim

Number of input dimensions to the MLP.

Type:

int

num_hidden_dim

Number of hidden dimensions in each hidden layer.

Type:

int

num_output_dim

Number of output dimensions from the MLP.

Type:

int

num_hidden_layers

Number of hidden layers in the MLP.

Type:

int

hidden_activation

Activation function for hidden layers. Choices are "identity", "relu", "leakyrelu", "elu", "sigmoid", "tanh", "sin", "softplus", "swish".

Type:

str

output_activation

Activation function for the output layer. Choices are "identity", "relu", "leakyrelu", "elu", "sigmoid", "tanh", "sin", "softplus", "swish".

Type:

str

class adeptml.configs.PhysicsConfig(forward_func: Callable, jacobian_func: Callable | None = None, use_vjp: bool = False, use_split_vjp: bool = False)

Configuration class for physics-related functions.

Three backpropagation modes are supported, selected via use_vjp and use_split_vjp:

Jacobian mode (default, use_vjp=False, use_split_vjp=False)

forward_func(x, *args) -> ndarray jacobian_func(x, *args) -> ndarray shape (batch, out, in)

backward forms the VJP by multiplying grad_output @ jacobian. Use when the full Jacobian is cheap or already available.

VJP mode (use_vjp=True, use_split_vjp=False)

forward_func(x, *args) -> ndarray jacobian_func(x, grad_output, *args) -> ndarray shape (batch, in)

backward calls jacobian_func with the upstream gradient directly, avoiding explicit Jacobian materialisation. Use when a manual VJP is cheaper than the full Jacobian.

Split-VJP mode (use_split_vjp=True)

forward_func(x, *args) -> (ndarray, pullback_fn)

forward_func returns both the output and a pullback closure (e.g. the result of jax.vjp). The closure is stored during the forward pass and called with grad_output during the backward pass without re-running any part of the physics. jacobian_func is not used and may be omitted.

This is the most memory-efficient mode when the physics solver is expensive and intermediates are already cached by the underlying framework (e.g. JAX checkpointing via jax.checkpoint).

Example usage with a JAX physics function:

from my_physics import physics_forward
# physics_forward(x, *args) -> (y, pullback_fn)

phy_cfg = PhysicsConfig(
    forward_func=physics_forward,
    use_split_vjp=True,
)

forward_func

Physics forward function. Signature depends on the chosen mode — see above.

Type:

Callable

jacobian_func

Jacobian or VJP function. Required for Jacobian and VJP modes; not used in split-VJP mode.

Type:

Callable, optional

use_vjp

Enable VJP mode. Ignored when use_split_vjp=True.

Type:

bool

use_split_vjp

Enable split-VJP mode. Takes precedence over use_vjp.

Type:

bool

Training Utils

Training utilities: batch step helper and full training loop.

adeptml.train_utils.train(model, train_loader, test_loader, optimizer, loss_fn, scheduler, filename, epochs, print_training_loss=True, save_frequency=50, grad_clip: float | None = None)

Full training loop with TensorBoard logging and checkpointing.

Parameters:

• model (HybridModel) – The hybrid model to train.

• train_loader (torch.utils.data.DataLoader) – DataLoader yielding (x, y) or (x, y, *args) batches.

• test_loader (torch.utils.data.DataLoader) – DataLoader yielding validation batches in the same format.

• optimizer (torch.optim.Optimizer) – Initialized optimizer.

• loss_fn (callable) – Loss function (y_pred, y_true) -> scalar tensor.

• scheduler (torch.optim.lr_scheduler) – Learning rate scheduler.

• filename (str) – Base directory name for saving runs and TensorBoard logs. A sub-directory run_N is created automatically (N auto-increments).

• epochs (int) – Number of training epochs.

• print_training_loss (bool) – Print epoch loss to stdout when True.

• save_frequency (int) – Save a model checkpoint every this many epochs.

• grad_clip (float, optional) – Maximum L2 norm for gradient clipping. None disables clipping.

Returns:

The trained model (same object, modified in place).

Return type:

HybridModel

adeptml.train_utils.train_step(model: HybridModel, optimizer: torch.optim.Optimizer, loss_fn, scheduler=None, grad_clip: float | None = None)

Build and return a single-batch training step closure.

Parameters:

• model (HybridModel) – The hybrid model to train.

• optimizer (torch.optim.Optimizer) – Initialized optimizer.

• loss_fn (callable) – Loss function (y_pred, y_true) -> scalar tensor.

• scheduler (torch.optim.lr_scheduler, optional) – Learning rate scheduler stepped after each batch.

• grad_clip (float, optional) – If set, gradients are clipped to this maximum L2 norm before the optimizer step. Useful for physics-informed training where gradient spikes can occur.

Returns:

_batch_step(x, y, args, test=False) -> float — runs one forward (and optionally backward) pass and returns the scalar loss.

Return type:

callable