Framework adapters

Native frameworks supply gradients; UnifiedIG supplies baseline handling, path integration, output conventions, and completeness checks. Install the framework appropriate for your model and device separately, or use the torch, jax, or tensorflow extras for their standard dependency.

PyTorch

Pass a torch.nn.Module directly. Inputs must be a single tensor per sample, with one scalar or output vector per sample. UnifiedIG respects the module’s device and floating dtype and restores its prior training/evaluation state.

import numpy as np
import torch

from cbaseline import background
import unifiedig as uig

torch.manual_seed(6)
model = torch.nn.Sequential(
    torch.nn.Linear(3, 8),
    torch.nn.Tanh(),
    torch.nn.Linear(8, 1),
)
data = torch.tensor([[0.5, -0.2, 0.8]])
reference = torch.randn(200, 3)
with torch.no_grad():
    reference_predictions = model(reference).numpy()[:, 0]
f0 = float(reference_predictions.mean())
bg = background(
    predictions=reference_predictions,
    f0=f0,
    features=reference.numpy(),
    weighting="calibrated",
)

# A deliberate single starting point is also valid: uig.Explainer(model, x0).
explanation = uig.Explainer(model, bg)(data)
print(explanation.values)
print(explanation.max_abs_completeness_error)

JAX

Wrap a differentiable prediction function in uig.JaxModel. The function is called with a batch by default. With explicit parameters, it receives (params, X); otherwise it receives X. Use vectorize=True for a function written for a single observation. Flax, NNX, Equinox, and Haiku models can be exposed through this calling convention; no additional UnifiedIG adapter is needed for each library.

import jax.numpy as jnp
import numpy as np

from cbaseline import background
import unifiedig as uig


def predict(params, X):
    hidden = jnp.tanh(X @ params["hidden_weights"])
    return hidden @ params["output_weights"]


params = {
    "hidden_weights": jnp.array([[0.8, -0.3], [0.2, 0.7]]),
    "output_weights": jnp.array([1.1, -0.6]),
}
rng = np.random.default_rng(4)
reference = rng.normal(size=(200, 2)).astype(np.float32)
reference_predictions = np.asarray(predict(params, jnp.asarray(reference)))
f0 = float(reference_predictions.mean())
bg = background(
    predictions=reference_predictions,
    f0=f0,
    features=reference,
    weighting="calibrated",
)
X_eval = np.array([[0.5, -0.2], [1.0, 0.4]], dtype=np.float32)

# JAX has no single fitted-model protocol: prediction functions may keep their
# parameters in a separate pytree, capture them in a closure, or belong to a
# framework such as Flax or Equinox. JaxModel is a lightweight adapter that
# records how UnifiedIG should call the function. It does not convert, copy,
# train, or otherwise modify the model or its parameters.
jax_model = uig.JaxModel(predict, params=params)

# A deliberate single starting point is also valid:
# uig.Explainer(jax_model, x0).
explanation = uig.Explainer(jax_model, bg)(X_eval)

print(explanation.values)
print(explanation.max_abs_completeness_error)

TensorFlow and Keras

TensorFlow-backed Keras models work directly. Wrap an arbitrary differentiable TensorFlow function with uig.TensorFlowModel. Keras 3 uses its configured TensorFlow, JAX, or PyTorch backend; configure it before importing Keras.

import numpy as np
from tensorflow import keras

from cbaseline import background
import unifiedig as uig

keras.utils.set_random_seed(9)
model = keras.Sequential(
    [
        keras.Input((3,)),
        keras.layers.Dense(8, activation="tanh"),
        keras.layers.Dense(1),
    ]
)
data = np.array([[0.5, -0.2, 0.8]], dtype=np.float32)
rng = np.random.default_rng(9)
reference = rng.normal(size=(200, 3)).astype(np.float32)
reference_predictions = model(reference, training=False).numpy()[:, 0]
f0 = float(reference_predictions.mean())
bg = background(
    predictions=reference_predictions,
    f0=f0,
    features=reference,
    weighting="calibrated",
)

# A deliberate single starting point is also valid: uig.Explainer(model, x0).
explanation = uig.Explainer(model, bg)(data)

print(explanation.values)
print(explanation.max_abs_completeness_error)

Output contract

Supply raw scores for classification. Visible sigmoid/softmax Keras heads are rejected, but UnifiedIG cannot inspect every transformation hidden in a custom function. Vector outputs default to class scores; use output_kind="regression" for multiple regression outputs. Multiple input tensors, dictionaries of outputs, and arbitrary structured outputs are outside the supported contract. See output interpretation.

The docs example checker runs framework examples when their dependencies are installed. Optional-framework CI jobs also check these examples alongside the backend tests; the documentation build itself needs only core dependencies.