Inference

Wrap a trained model as a Layer-1 Operator so it drops into a Sequential or Graph like any other step. ModelOp is framework- agnostic (duck-typed predict / __call__); SklearnModelOp and JaxModelOp add framework-specific adapters. The ML backends are imported lazily, so installing JAX or scikit-learn is only required for the wrapper you actually use.

ModelOp

ModelOp(model: Any, *, method: str = 'predict', feature_dim: str = 'feature', output_name: str = 'prediction', output_dim: str = 'output')

Bases: Operator

Wrap a trained model as an :class:~pipekit.Operator.

The wrapped model is expected to expose a method named method (default "predict") that maps a 2-D (n_samples, n_features) array to either a 1-D (n_samples,) or 2-D (n_samples, n_outputs) array. ModelOp reshapes leading non-feature dimensions of the input automatically.

Parameters:

Name	Type	Description	Default
model	Any	Any object exposing the configured method. Not stored in :meth:get_config; the reference is replaced by the literal string "<model>" so configs are JSON-safe.	required
method	str	Name of the model method to call. Defaults to "predict". Common alternatives are "predict_proba" and "transform".	'predict'
feature_dim	str	Name of the feature dimension on the input :class:xarray.DataArray. All other dims are treated as sample dims and preserved on the output.	'feature'
output_name	str	Name attached to the returned :class:xarray.DataArray. Defaults to "prediction".	'prediction'
output_dim	str	Name of the dimension along the model's output channels (only used when the model returns a 2-D array). Defaults to "output".	'output'

Notes

This class follows D4: no backend (sklearn / jax / torch) is imported here. The model is invoked via duck typing.

Source code in packages/xrtoolz-ops/src/xrops/inference/modelop.py

def __init__(
    self,
    model: Any,
    *,
    method: str = "predict",
    feature_dim: str = "feature",
    output_name: str = "prediction",
    output_dim: str = "output",
) -> None:
    self.model = model
    self.method = method
    self.feature_dim = feature_dim
    self.output_name = output_name
    self.output_dim = output_dim

get_config

get_config() -> dict[str, Any]

Return a JSON-serializable summary of the wrapper config.

The wrapped model is referenced as the literal string "<model>" because most trained models are not JSON-safe.

Source code in packages/xrtoolz-ops/src/xrops/inference/modelop.py

def get_config(self) -> dict[str, Any]:
    """Return a JSON-serializable summary of the wrapper config.

    The wrapped ``model`` is referenced as the literal string
    ``"<model>"`` because most trained models are not JSON-safe.
    """
    return {
        "model": "<model>",
        "method": self.method,
        "feature_dim": self.feature_dim,
        "output_name": self.output_name,
        "output_dim": self.output_dim,
    }

SklearnModelOp

SklearnModelOp(model: Any, *, method: str = 'predict', **kwargs: Any)

Bases: ModelOp

ModelOp flavored for scikit-learn estimators.

Functionally identical to :class:ModelOp, but documents the expected predict / predict_proba interface and validates, via duck typing, that the wrapped model exposes the requested method. sklearn itself is never imported.

Parameters:

Name	Type	Description	Default
model	Any	A fitted sklearn-style estimator.	required
method	str	Either "predict" (default) or "predict_proba"; any other method exposed by the estimator is also accepted.	'predict'
**kwargs	Any	Forwarded to :class:ModelOp.	{}

Source code in packages/xrtoolz-ops/src/xrops/inference/modelop.py

def __init__(
    self,
    model: Any,
    *,
    method: str = "predict",
    **kwargs: Any,
) -> None:
    if not hasattr(model, method):
        raise AttributeError(
            f"sklearn-style model {type(model).__name__!r} has no "
            f"method {method!r}; pass method= explicitly if needed."
        )
    super().__init__(model, method=method, **kwargs)

JaxModelOp

JaxModelOp(model: Any, *, method: str | None = None, jit: bool = True, **kwargs: Any)

Bases: ModelOp

ModelOp flavored for JAX / pytree-callable models.

The model can be any pytree-callable: an equinox.Module, a plain Python function closing over JAX arrays, a flax module paired with bound parameters, etc. By default the model is invoked as model(x); pass method= to dispatch to a named method.

When jit=True (default), the dispatch function is compiled with :func:jax.jit on first use. jax is imported lazily inside :meth:_resolve_fn.

Parameters:

Name	Type	Description	Default
model	Any	A callable / pytree-callable accepting a 2-D array.	required
method	str | None	Optional method name. None (default) means call the model itself as model(x).	None
jit	bool	If True, jit-compile the dispatch function. Defaults to True.	True
**kwargs	Any	Forwarded to :class:ModelOp (feature_dim, output_name, output_dim).	{}

Source code in packages/xrtoolz-ops/src/xrops/inference/modelop.py

def __init__(
    self,
    model: Any,
    *,
    method: str | None = None,
    jit: bool = True,
    **kwargs: Any,
) -> None:
    # Use a sentinel string for the parent's bookkeeping; we
    # override _resolve_fn so the actual method name is irrelevant
    # to the base class.
    super().__init__(model, method=method or "__call__", **kwargs)
    self._method_override = method
    self.jit = jit
    self._compiled: Callable[[np.ndarray], Any] | None = None