Inference ModelOp demo¶
This notebook shows how to wrap a trained sklearn estimator and a small JAX model as xrtoolz Layer-1 operators with ModelOp, then compose them in a Graph together with an RMSE metric.
Key points (per design decision D4):
xrtoolz.inference.ModelOpis framework-agnostic — it never importssklearn,jax,torch, orequinoxat module load time.SklearnModelOpandJaxModelOpare thin subclasses; backend imports happen lazily on first use.- Inputs are
xr.DataArray/xr.Datasetwith afeaturedim; outputs come back asxr.DataArraycarrying the input's non-feature coordinates.
In [ ]:
Copied!
import subprocess
import sys
try:
import google.colab # noqa: F401
IN_COLAB = True
except ImportError:
IN_COLAB = False
if IN_COLAB:
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
"-q",
"xrtoolz[inference] @ git+https://github.com/jejjohnson/xrtoolz@main",
],
check=True,
)
import subprocess
import sys
try:
import google.colab # noqa: F401
IN_COLAB = True
except ImportError:
IN_COLAB = False
if IN_COLAB:
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
"-q",
"xrtoolz[inference] @ git+https://github.com/jejjohnson/xrtoolz@main",
],
check=True,
)
In [ ]:
Copied!
import warnings
warnings.filterwarnings("ignore", message=r".*IProgress.*")
import warnings
warnings.filterwarnings("ignore", message=r".*IProgress.*")
In [ ]:
Copied!
import importlib.util
try:
from IPython import get_ipython
ipython = get_ipython()
except ImportError:
ipython = None
if ipython is not None and importlib.util.find_spec("watermark") is not None:
ipython.run_line_magic("load_ext", "watermark")
ipython.run_line_magic(
"watermark",
"-v -m -p numpy,xarray,scikit-learn,jax,equinox,xrtoolz",
)
else:
print("watermark extension not installed; skipping reproducibility readout.")
import importlib.util
try:
from IPython import get_ipython
ipython = get_ipython()
except ImportError:
ipython = None
if ipython is not None and importlib.util.find_spec("watermark") is not None:
ipython.run_line_magic("load_ext", "watermark")
ipython.run_line_magic(
"watermark",
"-v -m -p numpy,xarray,scikit-learn,jax,equinox,xrtoolz",
)
else:
print("watermark extension not installed; skipping reproducibility readout.")
In [1]:
Copied!
from __future__ import annotations
import numpy as np
import xarray as xr
rng = np.random.default_rng(0)
from __future__ import annotations
import numpy as np
import xarray as xr
rng = np.random.default_rng(0)
1. Synthetic data¶
Build a small regression dataset as an xr.DataArray with (sample, feature) dims and a matching target.
In [2]:
Copied!
n_samples, n_features = 256, 4
x = rng.standard_normal((n_samples, n_features))
true_w = np.array([1.5, -2.0, 0.5, 0.0])
y = x @ true_w + 0.1 * rng.standard_normal(n_samples)
X = xr.DataArray(
x,
dims=('sample', 'feature'),
coords={'sample': np.arange(n_samples),
'feature': ['x0', 'x1', 'x2', 'x3']},
name='X',
)
Y = xr.DataArray(y, dims=('sample',), coords={'sample': np.arange(n_samples)},
name='Y')
X
n_samples, n_features = 256, 4
x = rng.standard_normal((n_samples, n_features))
true_w = np.array([1.5, -2.0, 0.5, 0.0])
y = x @ true_w + 0.1 * rng.standard_normal(n_samples)
X = xr.DataArray(
x,
dims=('sample', 'feature'),
coords={'sample': np.arange(n_samples),
'feature': ['x0', 'x1', 'x2', 'x3']},
name='X',
)
Y = xr.DataArray(y, dims=('sample',), coords={'sample': np.arange(n_samples)},
name='Y')
X
Out[2]:
<xarray.DataArray 'X' (sample: 256, feature: 4)> Size: 8kB
array([[ 0.12573022, -0.13210486, 0.64042265, 0.10490012],
[-0.53566937, 0.36159505, 1.30400005, 0.94708096],
[-0.70373524, -1.26542147, -0.62327446, 0.04132598],
...,
[-0.03474712, -0.65237811, -1.0540028 , -0.66435631],
[ 1.07153834, 0.37416181, 0.58695536, 1.37996792],
[-1.17943093, 0.50995242, -1.07507411, -0.3343326 ]],
shape=(256, 4))
Coordinates:
* sample (sample) int64 2kB 0 1 2 3 4 5 6 7 ... 249 250 251 252 253 254 255
* feature (feature) <U2 32B 'x0' 'x1' 'x2' 'x3'2. Wrap a fitted scikit-learn model with SklearnModelOp¶
In [3]:
Copied!
from sklearn.linear_model import LinearRegression
from xrtoolz.inference import SklearnModelOp
lr = LinearRegression().fit(x, y)
sk_op = SklearnModelOp(lr)
sk_op
from sklearn.linear_model import LinearRegression
from xrtoolz.inference import SklearnModelOp
lr = LinearRegression().fit(x, y)
sk_op = SklearnModelOp(lr)
sk_op
Out[3]:
SklearnModelOp(model='<model>', method='predict', feature_dim='feature', output_name='prediction', output_dim='output')
In [4]:
Copied!
y_hat_lr = sk_op(X)
y_hat_lr
y_hat_lr = sk_op(X)
y_hat_lr
Out[4]:
<xarray.DataArray 'prediction' (sample: 256)> Size: 2kB
array([ 0.77974879, -0.8667726 , 1.17399827, -3.66898406, 0.02916046,
-3.27385027, 0.78937079, -1.63755358, -1.2149649 , -0.32181319,
-4.24246668, 1.76364757, 0.23971266, -1.97823141, -1.35275446,
2.58166046, 1.81393759, 5.40061897, 3.67217566, -2.84582535,
1.36620401, -2.88457878, -4.66964217, 0.95591679, 0.74101714,
-1.31546754, 3.51698973, -4.43719805, -0.16404806, 1.09521442,
-0.47415953, 0.62286438, 0.68961547, -2.12732199, -1.41281927,
-0.15390006, 0.23595025, -3.51766598, -0.24706437, 0.06783103,
1.09866068, 2.86683772, 1.16497463, 0.93956521, -5.25200251,
2.10175045, -1.94102501, 0.49652209, 1.58268359, -0.98244626,
-0.57505798, 2.25028592, -3.40823023, -2.44244715, 0.67171924,
3.60629567, -2.97063509, 0.96107722, -2.02629386, -2.28256376,
3.07125841, 2.03726277, 2.07158467, 3.85514171, -1.7517377 ,
-0.68214155, -1.85888621, 0.93982176, 1.02370607, 0.89314889,
1.31899618, -5.16567685, -3.01858091, -1.04388104, -1.68150174,
0.81332165, 0.36934065, -0.87634363, 4.58450973, -1.08770359,
1.08082191, 1.87591794, 0.8399639 , 1.61004792, 3.24215247,
-1.28547552, 2.39615948, 0.50195192, 1.9253275 , -0.56111335,
-5.58971532, 0.78679325, -2.586934 , -2.20042713, 0.08068475,
0.90436776, -0.70451924, -1.8360225 , 1.12784247, 4.7707012 ,
...
-1.7543918 , 3.75899693, -3.33481999, 1.44136993, 2.15520718,
-1.2180968 , -3.45013285, -2.25226173, 2.39995703, 0.20362997,
0.93997685, 3.18838279, -0.78114188, -4.22159983, -2.61501581,
-0.33910802, -0.67169236, -0.89909324, 3.60352299, -0.57624131,
0.21401295, 1.73643988, 4.64959104, -1.290047 , -1.65046671,
-4.75954152, 0.10747849, 2.16210019, 6.58496798, -3.82652097,
-3.29606287, 3.7866662 , -1.82259302, -5.73109074, -0.76041541,
1.0840726 , 0.45538886, 0.61371245, -1.65589355, 0.67945118,
-2.38360187, 0.7936815 , -1.53323094, 4.10089105, 2.58518433,
4.08037285, 0.53563313, 1.14942898, -1.38901549, -0.56542213,
-2.4664134 , -2.59714807, 4.05106964, -0.00966454, -0.26744949,
2.12321202, 0.67625738, 0.69617079, 2.58990524, 0.57169438,
-1.86321777, 0.28287701, 0.41352477, -0.84377892, 3.51755376,
-3.43618577, 3.62734658, -2.30547164, -1.56335068, 0.09466743,
0.95237963, -0.6325894 , 0.22929533, -3.75751754, 2.52756556,
3.10548238, 3.59717367, -3.45017701, -2.27099024, -0.92512039,
-2.73806306, -2.7682053 , -0.40369796, -1.35701526, 0.17083812,
2.56978281, 2.5818031 , 1.46314484, -1.08279843, -3.9363852 ,
1.25983003, -1.74603387, -0.76319116, 0.72658766, 1.14850786,
-3.33257957])
Coordinates:
* sample (sample) int64 2kB 0 1 2 3 4 5 6 7 ... 249 250 251 252 253 254 255Notice that the output is an xr.DataArray on the same sample axis as the input. The model itself was never serialized; get_config() captures only the wrapper config:
In [5]:
Copied!
sk_op.get_config()
sk_op.get_config()
Out[5]:
{'model': '<model>',
'method': 'predict',
'feature_dim': 'feature',
'output_name': 'prediction',
'output_dim': 'output'}
3. A tiny JAX model wrapped with JaxModelOp¶
We define a minimal MLP as a plain pytree-callable. JaxModelOp JIT-compiles it on first use; jax is imported lazily.
In [6]:
Copied!
import jax
import jax.numpy as jnp
from xrtoolz.inference import JaxModelOp
def init_mlp(key, in_dim, hidden, out_dim):
k1, k2 = jax.random.split(key)
return {
'W1': 0.1 * jax.random.normal(k1, (in_dim, hidden)),
'b1': jnp.zeros(hidden),
'W2': 0.1 * jax.random.normal(k2, (hidden, out_dim)),
'b2': jnp.zeros(out_dim),
}
params = init_mlp(jax.random.PRNGKey(0), n_features, 16, 1)
def mlp(params, x):
h = jnp.tanh(x @ params['W1'] + params['b1'])
return (h @ params['W2'] + params['b2']).squeeze(-1)
# Train for a few steps with plain SGD on the squared loss.
def loss(params, x, y):
pred = mlp(params, x)
return jnp.mean((pred - y) ** 2)
grad_fn = jax.jit(jax.grad(loss))
lr_rate = 0.05
for _ in range(200):
g = grad_fn(params, x, y)
params = jax.tree_util.tree_map(lambda p, gp: p - lr_rate * gp, params, g)
# Bind params and wrap. The wrapped callable takes a 2-D array.
def predict(x_arr):
return mlp(params, x_arr)
jax_op = JaxModelOp(predict, jit=True)
jax_op
import jax
import jax.numpy as jnp
from xrtoolz.inference import JaxModelOp
def init_mlp(key, in_dim, hidden, out_dim):
k1, k2 = jax.random.split(key)
return {
'W1': 0.1 * jax.random.normal(k1, (in_dim, hidden)),
'b1': jnp.zeros(hidden),
'W2': 0.1 * jax.random.normal(k2, (hidden, out_dim)),
'b2': jnp.zeros(out_dim),
}
params = init_mlp(jax.random.PRNGKey(0), n_features, 16, 1)
def mlp(params, x):
h = jnp.tanh(x @ params['W1'] + params['b1'])
return (h @ params['W2'] + params['b2']).squeeze(-1)
# Train for a few steps with plain SGD on the squared loss.
def loss(params, x, y):
pred = mlp(params, x)
return jnp.mean((pred - y) ** 2)
grad_fn = jax.jit(jax.grad(loss))
lr_rate = 0.05
for _ in range(200):
g = grad_fn(params, x, y)
params = jax.tree_util.tree_map(lambda p, gp: p - lr_rate * gp, params, g)
# Bind params and wrap. The wrapped callable takes a 2-D array.
def predict(x_arr):
return mlp(params, x_arr)
jax_op = JaxModelOp(predict, jit=True)
jax_op
Out[6]:
JaxModelOp(model='<model>', method=None, feature_dim='feature', output_name='prediction', output_dim='output', jit=True)
In [7]:
Copied!
y_hat_mlp = jax_op(X)
y_hat_mlp
y_hat_mlp = jax_op(X)
y_hat_mlp
Out[7]:
<xarray.DataArray 'prediction' (sample: 256)> Size: 1kB
array([ 0.8585566 , -0.9529432 , 1.2662219 , -3.6564898 , 0.0250868 ,
-3.3296552 , 0.8455371 , -1.7901497 , -1.3536094 , -0.3555792 ,
-4.0998425 , 1.9014722 , 0.24558553, -2.0843868 , -1.4654756 ,
2.6096785 , 1.9554522 , 4.9187512 , 3.5599341 , -2.9194057 ,
1.480791 , -2.9853766 , -4.380128 , 1.0449007 , 0.78173494,
-1.4467977 , 3.542401 , -4.250211 , -0.20193006, 1.2089403 ,
-0.5358871 , 0.66617143, 0.741186 , -2.2565196 , -1.349613 ,
-0.17765066, 0.24036115, -3.5870223 , -0.30619928, 0.03665807,
1.2006465 , 2.9162436 , 1.2640165 , 1.0330592 , -4.7423553 ,
2.2430131 , -2.0666053 , 0.5025524 , 1.6962843 , -1.0885069 ,
-0.6196619 , 2.3838575 , -3.3877306 , -2.6065967 , 0.6817821 ,
3.5966074 , -3.0897322 , 1.0268614 , -2.0732255 , -2.3054695 ,
3.1686711 , 2.0822198 , 2.2167494 , 3.868346 , -1.8111013 ,
-0.7731891 , -2.0182555 , 0.95451635, 1.0979984 , 0.94530743,
1.4413776 , -4.704519 , -3.1352308 , -1.118714 , -1.7825854 ,
0.80932844, 0.39608702, -0.9888877 , 4.404794 , -1.19252 ,
1.1662672 , 2.0287106 , 0.90040845, 1.7559241 , 3.2232695 ,
-1.3901957 , 2.5392666 , 0.49997202, 2.0218391 , -0.6267667 ,
-5.00229 , 0.8404076 , -2.7600827 , -2.2874897 , 0.09656513,
0.96243954, -0.80422324, -1.9963486 , 1.2339504 , 4.539151 ,
...
-1.932006 , 3.6789882 , -3.4057918 , 1.5564568 , 2.2922091 ,
-1.2432729 , -3.4910567 , -2.4315517 , 2.4760978 , 0.22106779,
0.9862792 , 3.2799897 , -0.8745411 , -4.085484 , -2.7747457 ,
-0.35869077, -0.76431364, -1.0161363 , 3.5707204 , -0.6370229 ,
0.22804126, 1.794624 , 4.4149055 , -1.4171939 , -1.8148192 ,
-4.453204 , 0.08648326, 2.3350468 , 5.537766 , -3.7391894 ,
-3.4061964 , 3.745376 , -1.988081 , -5.0936646 , -0.7808639 ,
1.1269488 , 0.48896775, 0.66723996, -1.801785 , 0.727231 ,
-2.557257 , 0.86809427, -1.6892625 , 4.0282803 , 2.7207677 ,
4.007917 , 0.57847023, 1.2637484 , -1.5372493 , -0.6340789 ,
-2.6445818 , -2.7442298 , 3.9474542 , -0.01326018, -0.2945381 ,
2.2346082 , 0.7486028 , 0.7595689 , 2.720043 , 0.6157526 ,
-2.0370483 , 0.30720034, 0.4487867 , -0.9328005 , 3.5776749 ,
-3.4659228 , 3.6655784 , -2.4403026 , -1.7277972 , 0.09813363,
1.0218886 , -0.70710236, 0.2164823 , -3.6706762 , 2.6547978 ,
3.2204788 , 3.5674105 , -3.4133852 , -2.3943906 , -1.0339239 ,
-2.7711997 , -2.9104075 , -0.4798171 , -1.490021 , 0.19324234,
2.5171888 , 2.7174227 , 1.5603739 , -1.1237477 , -3.9068513 ,
1.3811836 , -1.891633 , -0.86815995, 0.78171057, 1.2683421 ,
-3.411858 ], dtype=float32)
Coordinates:
* sample (sample) int64 2kB 0 1 2 3 4 5 6 7 ... 249 250 251 252 253 254 2554. Compose both wrappers in a single Graph with RMSE¶
Both ModelOp instances are full Layer-1 operators, so they slot into the functional Graph API alongside any other operator — here the RMSE metric from xrtoolz.metrics.
In [8]:
Copied!
from xrtoolz import Graph, Input, Operator
from xrtoolz.metrics.operators import RMSE
# RMSE is a two-input pixel metric over a dataset variable. Wrap each
# prediction + the reference into a small Dataset so it can be consumed.
class AsDataset(Operator):
"""Pack a prediction and reference DataArray into a Dataset."""
def __init__(self, name: str = 'y'):
self.name = name
def _apply(self, da):
return da.rename(self.name).to_dataset()
def get_config(self):
return {'name': self.name}
to_ds = AsDataset(name='y')
x_in = Input('X')
y_in = Input('Y')
rmse_op = RMSE(variable='y', dims='sample')
lr_pred = sk_op(x_in)
mlp_pred = jax_op(x_in)
lr_rmse = rmse_op(to_ds(lr_pred), to_ds(y_in))
mlp_rmse = rmse_op(to_ds(mlp_pred), to_ds(y_in))
graph = Graph(
inputs={'X': x_in, 'Y': y_in},
outputs={
'lr_pred': lr_pred,
'mlp_pred': mlp_pred,
'lr_rmse': lr_rmse,
'mlp_rmse': mlp_rmse,
},
)
print(graph.describe())
from xrtoolz import Graph, Input, Operator
from xrtoolz.metrics.operators import RMSE
# RMSE is a two-input pixel metric over a dataset variable. Wrap each
# prediction + the reference into a small Dataset so it can be consumed.
class AsDataset(Operator):
"""Pack a prediction and reference DataArray into a Dataset."""
def __init__(self, name: str = 'y'):
self.name = name
def _apply(self, da):
return da.rename(self.name).to_dataset()
def get_config(self):
return {'name': self.name}
to_ds = AsDataset(name='y')
x_in = Input('X')
y_in = Input('Y')
rmse_op = RMSE(variable='y', dims='sample')
lr_pred = sk_op(x_in)
mlp_pred = jax_op(x_in)
lr_rmse = rmse_op(to_ds(lr_pred), to_ds(y_in))
mlp_rmse = rmse_op(to_ds(mlp_pred), to_ds(y_in))
graph = Graph(
inputs={'X': x_in, 'Y': y_in},
outputs={
'lr_pred': lr_pred,
'mlp_pred': mlp_pred,
'lr_rmse': lr_rmse,
'mlp_rmse': mlp_rmse,
},
)
print(graph.describe())
Graph (2 inputs, 4 outputs): Inputs: ['X', 'Y'] [2] AsDataset(name='y') ← ['Y'] [3] AsDataset(name='y') ← ['Y'] [4] JaxModelOp(model='<model>', method=None, feature_dim='feature', output_name='prediction', output_dim='output', jit=True) ← ['X'] [5] SklearnModelOp(model='<model>', method='predict', feature_dim='feature', output_name='prediction', output_dim='output') ← ['X'] [6] AsDataset(name='y') ← ['<node 125852674600912>'] [7] AsDataset(name='y') ← ['<node 125849727031952>'] [8] RMSE(variable='y', dims='sample') ← ['<node 125849730559312>', '<node 125849727648096>'] [9] RMSE(variable='y', dims='sample') ← ['<node 125849726853008>', '<node 125849726933920>'] Outputs: ['lr_pred', 'mlp_pred', 'lr_rmse', 'mlp_rmse']
In [9]:
Copied!
results = graph(X=X, Y=Y)
{k: float(v.values) if v.shape == () else v for k, v in results.items() if 'rmse' in k}
results = graph(X=X, Y=Y)
{k: float(v.values) if v.shape == () else v for k, v in results.items() if 'rmse' in k}
Out[9]:
{'lr_rmse': 0.10134100923089427, 'mlp_rmse': 0.1696163804879099}
5. The wrapper is fully introspectable¶
Both wrappers expose JSON-serializable configs (the model itself is referenced as the literal string "<model>"):
In [10]:
Copied!
import json
print(json.dumps([sk_op.get_config(), jax_op.get_config()], indent=2))
import json
print(json.dumps([sk_op.get_config(), jax_op.get_config()], indent=2))
[
{
"model": "<model>",
"method": "predict",
"feature_dim": "feature",
"output_name": "prediction",
"output_dim": "output"
},
{
"model": "<model>",
"method": null,
"feature_dim": "feature",
"output_name": "prediction",
"output_dim": "output",
"jit": true
}
]