Mechanism 1: kinks (ReLU)

The rectified linear unit is made of two straight lines. Its second derivative is zero everywhere except at , where all of its curvature is concentrated in a kink (formally, is a delta function there). A ReLU network is therefore piecewise linear. In one dimension it is a polyline, and each neuron contributes one kink at whose size is the change in slope there.

Here is the simplest curved function, , built from ReLUs. Each unit adds a kink, and the polyline gets closer to the parabola as the kinks get closer together.

Top: x² (thick grey) and a sum of ReLU units (dark), with the individual weighted units in color. Bottom: the slopes. The slope of x² changes smoothly; the network’s slope is a staircase, and each step is one kink. Move the slider to add units.

With units, and so evenly spaced pieces on , the largest error is . The error falls quickly, but the curvature is never smooth: the slope changes only in jumps, and between the jumps .

Once you can make a square, you can make a product, because of the polarization identity

The right side is a difference of two squares of linear combinations of the inputs. That is exactly what neurons can see. Neurons with weights proportional to build the first square as a function of , neurons with weights proportional to build the second, and the output layer subtracts them. The network never multiplies by . It approximates two one-dimensional parabolas along two diagonal directions, and their difference happens to be a product.

This should give you pause about what a neural network “learns”. It is learning some equivalent looking representation of a function that goes through the data.

Mechanism 2: curvature (tanh, sigmoid, softplus, GELU)

Smooth activations like bend everywhere, so is nonzero almost everywhere. How much a neuron bends depends on where it operates on the curve, and the bias sets that operating point. Expand around an operating point :

At the quadratic term vanishes (tanh is odd, so it is locally a straight line with a cubic correction). On the shoulders of the curve the quadratic term is largest; peaks near . Now let be the neuron’s input. Then

and there is the product. The curvature of the activation turns a sum into a square, and the square of a sum contains a cross term. It always comes bundled with and terms, though, and with all the higher order terms of the series. To isolate the cross term we use the polarization identity again. With , four neurons give

Adding and cancels the odd terms (including the linear one), and subtracting the two pairs cancels the squares, leaving plus terms of order .

Left: tanh, the operating point set by the bias, and the local parabola there. The curvature is what makes products possible. Right: four tanh neurons combined as above, compared with the hyperbolic contours of x₁x₂ (green dashes). Small ε is accurate but needs enormous output weights.

There is a catch, and the right panel shows it. The approximation is only good when is small, so that each neuron stays in the region where the parabola describes it. But the output weights are , which become huge. The network makes a product by taking the difference of four large, nearly equal numbers. Here is the same thing in code.

def g(s, eps, b=0.66):
    return np.tanh(b + eps * s)

def tanh2(b):
    "Second derivative of tanh."
    t = np.tanh(b)
    return -2 * t * (1 - t**2)

grid = np.linspace(-1, 1, 101)
x1, x2 = np.meshgrid(grid, grid)
u, v = x1 + x2, x1 - x2
for eps in [1.0, 0.3, 0.1, 0.03]:
    c = 1 / (4 * eps**2 * tanh2(0.66))
    approx = c * (g(u, eps) + g(-u, eps)
                  - g(v, eps) - g(-v, eps))
    err = np.abs(approx - x1 * x2).max()
    print(f"eps = {eps:4.2f}  |output weight| = {abs(c):6.1f}"
          f"  max error = {err:.1e}")

eps = 1.00  |output weight| =    0.3  max error = 6.6e-01
eps = 0.30  |output weight| =    3.6  max error = 1.1e-01
eps = 0.10  |output weight| =   32.5  max error = 1.3e-02
eps = 0.03  |output weight| =  360.8  max error = 1.2e-03

This is typical of what a trained network does too: hidden units with large outputs of opposite sign that mostly cancel. We saw that cancellation in the learned basis functions of the last post.

Ratios are harder than products. Over a range where stays away from zero, is just another smooth curved function of one linear combination, and it can be approximated with kinks or curvature like was. Near it has a pole, and units are bounded and ReLU units grow only linearly, so no finite network gets close there. Then the ratio has to be assembled from that and a product, so it costs at least two of the tricks above.

One activation function would make all of this easy. If the inputs were and and the activation were , a single neuron would compute , with products, ratios and powers all exact. We will come back to that at the end, because it is really a statement about features.

Computing the coefficients

In general, if we expand in some basis ,

we find the coefficients by minimizing the squared error . Setting the derivative with respect to each to zero gives the normal equations

where is the inner product, is the Gram matrix, and . For an arbitrary basis this is a linear regression problem: every coefficient depends on every other one through .

The sines and cosines are special because they are orthogonal on :

That makes diagonal, and the normal equations decouple into one independent equation per coefficient, . For the Fourier series that is

No regression is required. Each coefficient is a projection of onto one basis function, and it is computed with an integral. There is a second benefit hiding in that formula: does not depend on . If we add more terms to the series, the coefficients we already have do not change. We just compute the new ones.

An example

Here is a lumpy function that is not obviously made of sines and cosines.

import numpy as np
from scipy.integrate import quad

def f(x):
    return np.exp(-2 * (x - 1)**2) + 0.6 * np.exp(-4 * (x + 1.5)**2)

def a(n):
    return quad(lambda x: f(x) * np.cos(n * x), -np.pi, np.pi)[0] / np.pi

def b(n):
    return quad(lambda x: f(x) * np.sin(n * x), -np.pi, np.pi)[0] / np.pi

def fourier(x, N):
    y = a(0) / 2 * np.ones_like(x)
    for n in range(1, N + 1):
        y += a(n) * np.cos(n * x) + b(n) * np.sin(n * x)
    return y

Top: f(x) and its Fourier partial sum. Bottom: the coefficients, each one computed from a single integral. Move the slider and the faded bars join the sum without any of the bars changing height.

With eight terms the series is indistinguishable from the function (the largest error is about 0.0015). It converges this quickly partly because almost vanishes at , so its periodic extension is smooth. A function that does not match up at the ends would show Gibbs ringing near them and converge much more slowly. Every coefficient came from a single integral, and adding a term just adds a bar to the chart; it never changes the ones already there.

Adding terms does not change the coefficients

To make the second point concrete, let’s fit the same function by least squares on a grid of points, first with 3 harmonics and then with 6. We do the same with a polynomial basis , which is not orthogonal.

xg = np.linspace(-np.pi, np.pi, 2001)[:-1]   # one period, evenly spaced
yg = f(xg)

def fourier_basis(x, N):
    cols = [np.ones_like(x)]
    for n in range(1, N + 1):
        cols += [np.cos(n * x), np.sin(n * x)]
    return np.column_stack(cols)

def poly_basis(x, N):
    return np.column_stack([x**k for k in range(N + 1)])

for name, basis, Ns in [("Fourier", fourier_basis, (3, 6)),
                        ("polynomial", poly_basis, (3, 6))]:
    print(name)
    for N in Ns:
        c = np.linalg.lstsq(basis(xg, N), yg, rcond=None)[0]
        print(f"  N = {N}: first four coefficients", np.round(c[:4], 4))

Fourier
  N = 3: first four coefficients [ 0.2841  0.2015  0.1376 -0.2312]
  N = 6: first four coefficients [ 0.2841  0.2015  0.1376 -0.2312]
polynomial
  N = 3: first four coefficients [ 0.4682  0.1336 -0.056  -0.0188]
  N = 6: first four coefficients [ 0.2641  0.3645  0.3232 -0.1281]

The first four Fourier coefficients are the same whether we keep 3 harmonics or 6 (and they are the projections , , , from before). The first four polynomial coefficients all change when we add , because those functions overlap with the ones already in the expansion. With a non-orthogonal basis, every time you add a function you have to solve for all the coefficients again.

One caveat: the grid matters. The sines and cosines are orthogonal as integrals, and they stay orthogonal as sums only on a uniform grid that covers exactly one period. That is why the last point is dropped with [:-1], since and are the same point on the circle. On that grid the largest off-diagonal entry of is below . On a non-uniform or non-periodic grid the discrete Fourier basis is not orthogonal, and the least squares coefficients would shift when you add terms too.

The Fourier series as a network

Now let’s draw the Fourier series the way we usually draw a neural network. There is one input, . It feeds a layer of “hidden units”, each of which computes or . The outputs of those units are multiplied by the coefficients and summed to give .

A truncated Fourier series drawn as a neural network with one hidden layer. The input weights are fixed integer frequencies; the output weights are the Fourier coefficients, drawn with a width proportional to their size.

It looks a lot like a neural network, and it really is one. Since , every hidden unit computes

with weight and bias . The output is a linear combination . That is exactly a one-hidden-layer network with a cosine activation function. It is a special one in two ways:

1. The first-layer weights and biases are fixed. We chose them ahead of time (integer frequencies, two phases) rather than learning them.
2. Because of that choice the hidden units are orthogonal, so the output weights can be computed one at a time by projection, and they never change when we add more hidden units.

A typical neural network uses something like or ReLU instead of . The logistic sigmoid, and later , came from the picture of a neuron firing: each is a smooth step. ReLU is a more modern choice, made because it makes deep networks easier to optimize. We also give up the integer frequencies of the Fourier basis. This gives us even more flexibility, but it comes at a cost: the hidden units are no longer orthogonal, and the first-layer weights now have to be fitted instead of computed. We have to find all the weights together by regression, and when we add more hidden units the old weights change.

There is a middle ground that makes the point even more clearly. In random-feature models and extreme learning machines, the first-layer weights are drawn at random and then frozen, and only the output weights are fitted. It is the same picture as the Fourier series, with a random basis in place of an orthogonal one. Because the basis is not orthogonal, the output weights come from least squares instead of projection.

The output is linear in the new space

The point of the transformation is that a function which is nonlinear in becomes linear in . It is the same trick as fitting : the model is nonlinear in , but it is linear in the feature , so we compute that feature and fit a line. (Linearizing the Arrhenius equation, , is a close cousin. There we transform the output to as well, while a network only transforms the input.)

We can see this geometrically with only one harmonic, so the feature space is small enough to draw. Take . As a function of it is a wave. In the feature space , the input interval is mapped onto a circle, and the data lie exactly on a flat plane above that circle.

Left: the data are nonlinear in x. Right: the same data in the feature space (cos x, sin x) lie on a plane, so finding the coefficients is linear regression. Hover a point to find it in both panels.

The grey circle on the floor is where the inputs land in feature space. The points sit on the plane, and fitting a plane is linear regression.

Deep networks

So far there is one hidden layer. What changes when we stack more of them? Let’s write a network with two hidden layers using matrices, for a batch of inputs ():

Each hidden layer is a linear map to a new dimension followed by an elementwise nonlinearity, just like Step 1 and Step 2 above. Stacking them is function composition. Call the layer maps and . The features at the end are

and the output is still linear in them,

However many layers there are, the last one is a linear model on the features, and with a squared-error loss, fitting it is linear regression. The columns of are basis functions, each one a composition of the layers before it. The only difference from the Fourier series is that these basis functions are not chosen in advance; their shapes are set by , which are fitted along with .

A network with two hidden layers in graph form (top) and matrix form (bottom). The hidden layers transform n inputs into n rows of 3 features, the columns of H₂. The output layer is linear regression on them. (Three features keep the drawing small; the trained example below uses six.)

Checking that the last layer is linear regression

Let’s train a small network on the function from Part 1 and look inside it. This one has two hidden layers of six units. To keep it to numpy and scipy, we pack all the weights into one vector and minimize the mean squared error with BFGS.

from scipy.optimize import minimize

sizes = [1, 6, 6]   # input, hidden layer 1, hidden layer 2
shapes = []
for m, k in zip(sizes[:-1], sizes[1:]):
    shapes += [(m, k), (k,)]      # W, b for each hidden layer
shapes += [(sizes[-1],), ()]      # c, c0 for the output layer

def unpack(p):
    out, i = [], 0
    for s in shapes:
        k = int(np.prod(s))
        out.append(p[i:i + k].reshape(s))
        i += k
    return out

def features(p, X):
    "The hidden layers: H = tanh(... tanh(X W1 + b1) ... W + b)."
    *hidden, c, c0 = unpack(p)
    H = X
    for W, b in zip(hidden[::2], hidden[1::2]):
        H = np.tanh(H @ W + b)
    return H

def nn(p, X):
    *_, c, c0 = unpack(p)
    return features(p, X) @ c + c0

xt = np.linspace(-np.pi, np.pi, 200)
Xt, yt = xt[:, None], f(xt)

def loss(p):
    return np.mean((nn(p, Xt) - yt)**2)

p0 = np.random.default_rng(42).normal(size=sum(int(np.prod(s)) for s in shapes))
sol = minimize(loss, p0, method="BFGS", options={"maxiter": 5000})
print(f"{len(p0)} parameters, MSE = {sol.fun:.2e}")

61 parameters, MSE = 4.81e-06

Now take the trained hidden layers, compute the features for the training inputs, and do ordinary linear regression of on them. If the last layer is linear regression, we should get back the output weights the network learned.

H2 = features(sol.x, Xt)
A = np.column_stack([H2, np.ones(len(xt))])
c_lstsq = np.linalg.lstsq(A, yt, rcond=None)[0]

*_, c_nn, c0_nn = unpack(sol.x)
print("network:", np.round(np.append(c_nn, c0_nn), 3))
print("lstsq:  ", np.round(c_lstsq, 3))
print(f"cond(A) = {np.linalg.cond(A):.0f}")

network: [-1.044  0.962 -1.522  1.387 -0.859  0.774 -0.994]
lstsq:   [-1.044  0.962 -1.523  1.387 -0.859  0.776 -0.994]
cond(A) = 159

If you run this yourself, expect different numbers. The path BFGS takes depends on the numpy, scipy and BLAS versions, so the trained weights will differ. What should reproduce is that the two rows agree.

They agree to within the tolerance the optimizer stopped at, and that is not a coincidence. Write and . With the hidden weights held fixed, the loss is a convex quadratic in . Its gradient, , vanishes at a stationary point of the full training problem, and that condition is exactly the normal equations. As long as has full rank (its condition number is modest here), the network’s output weights are the least squares solution on its own features.

Finally, here are the six basis functions the network built, the columns of . Compare them to the sines and cosines in Part 1: these are not orthogonal and they are not pretty, but a weighted sum of them reproduces .

Left: the columns of H₂, the basis functions learned by the hidden layers. Switch to cₖ·hₖ to see them scaled by their output weights. Right: their linear combination. Toggle basis functions off to see that the fit needs all of them.

So a neural network does the same two things as the Fourier series: it transforms the input into a new vector space with a nonlinear function, and it does linear regression in that space. The Fourier series fixes the transformation in advance and gets the regression for free. A neural network learns the transformation, and pays for it with a nonlinear optimization over all the weights at once.

You can see from the learned basis functions that there is a lot of cancellation: each value of the output is a sum of positive and negative contributions from the hidden units. That is a common property of learned features. They are not orthogonal, and they are rarely interpretable on their own; individually they mean little, and only their linear combination reproduces the output.

We would never ask “What did the Fourier series learn about our data?” The Fourier basis is fixed, and all the data-dependence is in the coefficients. A neural network is different in exactly one place: the features are learned too, so that is where most of what it “learned” lives. That is also why transfer learning works, since features learned on one problem are often a good basis for another. So these viewpoints do not make a network fully transparent, but they do make its structure transparent. It is a learned nonlinear feature transform followed by a linear model. The structure is simple; the learned features usually are not.

We used a simple dense network with a squared-error loss here, which makes the last layer linear regression. The same reading applies to other heads with a small change. A classifier with a softmax output is doing multinomial logistic regression on its features, so in general the last layer is a generalized linear model. It also applies to convolutional networks, recurrent networks, transformers, graph neural networks and many other architectures. What differs is how the features are generated.

Symbolic regression

SymbolicRegressor in jaxsr builds a library of candidate basis functions of your inputs (powers, logs, exponentials, products, and any you add), then selects a sparse subset and fits the coefficients. The result is an equation of the form y = Σ βᵢ φᵢ(x) that you can read, differentiate, and write straight into an optimization model. Because the final fit is linear in the coefficients, the uncertainty is well grounded. You get prediction intervals from least squares, and also Pareto-front ensembles, conformal intervals and residual bootstraps when you trust the model form less. You can add domain knowledge as constraints: bounds, signs on coefficients, monotonicity, convexity. These are enforced at the training points and, if you ask, on a grid, which is not the same as a guarantee everywhere. MultiOutputSymbolicRegressor fits several targets at once.

Examples: basic usage, uncertainty quantification, constraint enforcement, multi-output, chemical kinetics, ODE discovery, and the comprehensive tutorial.

Additive symbolic models

Sometimes the best description is a sum of many small, simple terms rather than one clever expression. The stagewise, backfitting and recursive regressors in jaxsr build that sum one piece at a time. They do not report uncertainty yet. Example: additive symbolic regression.

Superposition

Many experiments produce a family of curves, one per temperature or concentration, and the claim is that shifting the axes collapses them onto one master curve. That is time-temperature superposition in rheology, Larson-Miller curves in creep, isoconversional methods in kinetics. Usually you shift the curves by hand and then assume a form for the shift, such as Arrhenius or WLF. SuperpositionRegressor learns the shift law symbolically from a library of candidates, which lets you extrapolate to conditions you never measured. It gives intervals on the shift-law parameters, such as an activation energy; the predictions are point values. Example: superposition master curves.

Response surfaces

ResponseSurface is classical response-surface methodology on top of the symbolic regressor: design generators, coded variables, ANOVA, and canonical analysis of the stationary point. It also lets the model leave the quadratic when the data clearly want something else. If you are running a designed experiment with a handful of factors, start here. Examples: RSM formulation, DOE study workflow.

Model averaging

When several equations fit about equally well, picking one hides real uncertainty about the model form. BayesianModelAverage keeps them all, weighted by how well the data support each one, so the disagreement between plausible models shows up in the prediction. Examples: uncertainty quantification, comparing isotherm models.

Symbolic classifier

SymbolicClassifier applies the same sparse selection to a logistic model, fitted by iteratively reweighted least squares, with one-vs-rest for more than two classes. You get a decision boundary you can write down. Example: symbolic classification.

iBART

jax-ibart is a JAX version of iBART, iterative Bayesian additive regression trees, from Liu et al. (J. Chem. Phys. 156, 164105, 2022). It goes after the same thing as SISSO, a compact descriptor built from the inputs, but it never writes down the whole candidate space. Each round applies the operators (products, ratios, powers, logs, exponentials and so on) to the features that survived the last round, and BART screening keeps only the few that matter. That keeps the pool small, where full enumeration can reach 10¹⁷ candidates. A sparse linear model is fitted on the survivors at the end. Every descriptor keeps its SymPy expression, and a unit filter throws out ones like volume + speed before they are ever screened. It is not on PyPI; install it from GitHub.

Selection is the largest source of uncertainty here, because a different sample of data can pick a different descriptor. jax_ibart.uncertainty has split-conformal intervals that wrap the whole fit, selection included, and selection-stability counts across bootstrap resamples, which tell you whether the data pin the descriptor down. It also has OLS intervals on the final model, which are optimistic because they condition on the selection, and the BART posterior. These are functions you call, not predict(X, return_std=True).

Because the model is algebra, model.to_discopt(m) writes it into a discopt problem and you can optimize over it globally with a certificate. Bounds on the inputs default to the training range, and guards keep the solver away from poles and logs of negative numbers that the data never visited. Examples: quickstart, single-atom catalysis, reaction kinetics, uncertainty, optimizing over a model.

SISSO

SISSO (sure independence screening and sparsifying operator) takes a different route to the same goal. It composes the inputs into a very large space of candidate descriptors, screens them, and keeps the few that explain the data. It is popular in materials science for finding descriptors. The pycse wrapper sits on TorchSISSO. SISSO gives hat-matrix intervals for one equation, and SISSOEnsemble trains a shallow ensemble of equations with an NLL or CRPS loss. Examples: SISSO, SISSO ensemble.

regress and nlinfit

Do not skip the simplest option. If you know the form of the model, pycse.regress and nlinfit fit it by linear or nonlinear least squares and give confidence intervals on the parameters. predict and nlpredict carry those intervals through to the predictions. With few data points this is often the most honest model you can have. LinearRegressionUQ wraps the linear case as an sklearn estimator. Example: regression functions.

Linear and hyperplane trees

jax-ldt grows decision trees with a linear model in each leaf. LinearTreeRegressor splits on one input at a time. HyperplaneTreeRegressor splits on oblique hyperplanes w·x ≤ b, so it needs far fewer leaves when the real boundaries are not aligned with an axis. The model is piecewise linear, which is the point: it embeds exactly in a mixed-integer linear program, for example in discopt, and the solver then optimizes the surrogate globally. Uncertainty comes from each leaf’s parameter covariance, propagated linearly or to second order, or from conformal intervals computed per leaf. Trees also export to ONNX. Examples: quickstart, uncertainty, active learning, MIP optimization, ONNX export, steam tables.

Differentiable trees

jax-difftree makes the splits soft, so the whole tree trains end to end by gradient descent. Each input follows every path with some probability. Each leaf holds a Gaussian, multivariate if there are several outputs, so a prediction is a mixture of Gaussians. That gives a variance that changes with x, and it splits into a within-leaf (aleatoric) part and a between-leaf (routing) part. With several outputs you can condition on the ones you have measured to predict the rest, or ask which x would most likely give a target y. For optimization, affine leaves can be hardened into a crisp piecewise-affine model and exported for a MILP, so this sits between the smooth world and the mixed-integer world. DifferentiableTreeClassifier splits its predictive entropy into aleatoric and epistemic parts. It does not have a notebook yet. Examples: heteroscedastic regression, multivariate inverse design, piecewise-linear MILP, failure modes, scaling.

Leaf-model trees

LeafModelRegressor in pycse is the simplest version of the idea. An ordinary sklearn decision tree splits the data, then any sklearn model you like is fitted in each leaf, and its uncertainty is used if it has one. By default the splits are chosen as if the leaves were constant. refine_splits=True re-chooses them for the leaf model, in the style of M5 model trees. I think of it as a proof of concept and a quick baseline next to jax-ldt. Example: leaf models.

Convex neural networks (ICNN)

An input-convex network keeps the weights on its hidden path non-negative and uses convex, non-decreasing activations, which makes f convex in x. If you minimize the surrogate, any local minimum you find is the global one. The demo also shows how to use one as a convex lower bound of a non-convex function. JAXICNNRegressor gives point predictions only, so pair it with another model when you need intervals.

Monotonic and periodic networks

JAXMonotonicRegressor is guaranteed increasing or decreasing in the inputs you name. Use it when you know the sign of an effect, such as a rate that rises with temperature, and do not want the surrogate to wiggle against it where data are sparse. JAXPeriodicRegressor puts Fourier features on chosen inputs so the model is exactly periodic in them. Both get uncertainty from last-layer prediction rigidity (see LLPR). Examples: monotonic, periodic.

KANs

Kolmogorov-Arnold networks put learnable one-dimensional functions, here B-splines, on the edges instead of fixed activations on the nodes. You can plot each edge function, which makes them easier to inspect than an MLP. KAN gets uncertainty from an ensemble of output heads calibrated on a validation set. KANLLPR gets it from prediction rigidity instead. Examples: KAN, KAN + LLPR.

ZENN

The zentropy-enhanced neural network is a thermodynamics-inspired model for heterogeneous data from several sources. It predicts an energy at a temperature and can reconstruct the energy landscape. It reports epistemic and aleatoric uncertainty, and the NLL variant calibrates the aleatoric part to the noise. It is described in PNAS 123(1), e2511227122. The code is in pycse, but it has no example notebook yet.

LLPR

Last-layer prediction rigidity is the cheapest way I know to add uncertainty to a network you have already trained. It uses the last-layer features and a covariance computed from them, with no retraining and no ensemble, and is calibrated on held-out data. LLPRRegressor is the standalone version, and the same idea gives the uncertainty in the monotonic, periodic and KAN-LLPR models. Examples: LLPR demo, LLPR study.

DPOSE

DPOSE (direct propagation of shallow ensembles) shares the whole network across ensemble members except the last layer, and trains with a CRPS or NLL loss so the spread is calibrated. It costs little more than a single network. Because you have the members, you can push any function of the prediction through them and get its uncertainty directly. Example: DPOSE.

k-fold networks

KfoldNN trains each output neuron on a different fold of the data, so one network carries a k-member ensemble, and the spread across outputs is the uncertainty. Example: k-fold NN.

NN + Bayesian linear and NN + GMM

These two use an sklearn MLP as a feature generator and put a probabilistic model on top. NeuralNetworkBLR uses Bayesian linear regression as the last layer. NeuralNetworkGMM uses Gaussian mixture regression, which can represent noise that changes with x and predictions with more than one mode. They are not JAX models, so they do not fit the differentiable story as neatly, but they are simple and robust. Examples: NN + Bayesian linear, NN + GMM.

Normalizing flows

NFlowsRegressor learns the whole conditional density p(y | x) as an invertible transformation of a simple distribution. You can sample from it and read off quantiles. Because the transformation is invertible, it also runs backwards to infer which inputs are consistent with an observed output. It is built on PyTorch (nflows), not JAX. Example: normalizing flow regression, which includes a CSTR design problem.

Active learning

Uncertainty is not only for error bars. It tells you where to sample next. Any model above that returns a standard deviation can drive an active-learning loop that adds data where the surrogate is least sure. Examples: jaxsr active learning and Langmuir DOE, jax-ldt active learning, pycse active learning and active surrogates.

Comparing uncertainty

The methods do not agree, and they fail in different ways. The pycse UQ comparison runs conformal intervals (MAPIE), NN + Bayesian linear, NN + GMM and DPOSE on the same data and checks coverage and heteroscedastic noise. The jax-difftree failure modes notebook shows where soft trees go wrong.