Skip to main content
API REFERENCE/omnibias-torch

omnibias.torch.architectures.joint_operator

0.5.0rc1Release candidatePython >=3.10Apache-2.0

PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.

import omnibias.torch.architectures.joint_operator
Module contract & mathematical conventions

Joint operator-predictor learning for tabular scientific regression.

The model is a first practical step toward the "Omnibias as scientific AutoML/foundation tool" direction: build a named bank of differentiable operators, learn sparse gates over that bank, and train the predictor head in the same optimization loop.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →
CLASS

FittedJointOperatorRegressor​

omnibias.torch.architectures.FittedJointOperatorRegressorSource ↗
class FittedJointOperatorRegressor

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

A fitted joint operator model plus normalization state.

Parameters

NameType & defaultContract
modeldataclass fieldJointOperatorRegressorRequiredSee the operation contract and implementation below.
x_meandataclass fieldNDArray[np.float32]RequiredSee the operation contract and implementation below.
x_scaledataclass fieldNDArray[np.float32]RequiredSee the operation contract and implementation below.
y_meandataclass fieldfloatRequiredSee the operation contract and implementation below.
y_scaledataclass fieldfloatRequiredSee the operation contract and implementation below.
historydataclass fielddict[str, list[float]]RequiredSee the operation contract and implementation below.
Executed regression · joint operator fit accepts sample weights

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_accepts_sample_weights() -> None:
    x, y = _synthetic(n=500)
    weights = np.ones(350, dtype=np.float32)
    weights[y[:350] < np.median(y[:350])] = 2.0
    fitted = fit_joint_operator_regressor(
        x[:350],
        y[:350],
        x[350:],
        y[350:],
        seed=1,
        epochs=80,
        patience=30,
        train_sample_weight=weights,
        val_sample_weight=np.ones(150, dtype=np.float32),
        model_kwargs={"ombu_channels": 0, "stochastic_gates": False},
    )
    pred = fitted.predict(x[350:])
    assert np.isfinite(pred).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_fit_accepts_sample_weights" -q
Regression example · joint operator fit accepts asymmetric loss

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_accepts_asymmetric_loss() -> None:
    x, y = _synthetic(n=420)
    fitted = fit_joint_operator_regressor(
        x[:300],
        y[:300],
        x[300:],
        y[300:],
        seed=2,
        epochs=60,
        patience=25,
        asymmetric_weight=1e-3,
        model_kwargs={"ombu_channels": 0, "stochastic_gates": False},
    )
    assert "val_asymmetric" in fitted.history
    assert np.isfinite(fitted.history["val_asymmetric"]).all()
python -m pytest packages/omnibias-torch/tests/test_joint_operator.py -k test_joint_operator_fit_accepts_asymmetric_loss -q

Fields & methods

modelfield
model: JointOperatorRegressor

Declared JointOperatorRegressor field on FittedJointOperatorRegressor. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

Regression example · joint operator fit is robust to float64 default dtype

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_is_robust_to_float64_default_dtype() -> None:
    """A process-wide ``float64`` default must not desync the float32 model
    from its float32 inputs.

    Regression for the ``addmv ... Float, Double, Float`` crash in the ridge
    polish (and the matching nn.Linear mismatch on the OMBU bank path) that
    triggered when a user ran under ``torch.set_default_dtype(torch.float64)``.
    The autouse conftest fixture restores the default after this test.
    """
    torch.set_default_dtype(torch.float64)
    x, y = _synthetic(n=240)
    fitted = fit_joint_operator_regressor(
        x[:160],
        y[:160],
        x[160:],
        y[160:],
        seed=0,
        epochs=40,
        patience=20,
        model_kwargs={"ombu_channels": 2, "stochastic_gates": False},
    )
    assert fitted.model.readout.weight.dtype == torch.float32
    assert fitted.model.gate_logits.dtype == torch.float32
    pred = fitted.predict(x[160:])
    assert np.isfinite(pred).all()
python -m pytest packages/omnibias-torch/tests/test_joint_operator.py -k test_joint_operator_fit_is_robust_to_float64_default_dtype -q
View source ↗
x_meanfield
x_mean: NDArray[np.float32]

Declared NDArray[np.float32] field on FittedJointOperatorRegressor. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
x_scalefield
x_scale: NDArray[np.float32]

Declared NDArray[np.float32] field on FittedJointOperatorRegressor. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
y_meanfield
y_mean: float

Declared float field on FittedJointOperatorRegressor. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
y_scalefield
y_scale: float

Declared float field on FittedJointOperatorRegressor. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
historyfield
history: dict[str, list[float]]

Declared dict[str, list[float]] field on FittedJointOperatorRegressor. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

Regression example · joint operator fit accepts asymmetric loss

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_accepts_asymmetric_loss() -> None:
    x, y = _synthetic(n=420)
    fitted = fit_joint_operator_regressor(
        x[:300],
        y[:300],
        x[300:],
        y[300:],
        seed=2,
        epochs=60,
        patience=25,
        asymmetric_weight=1e-3,
        model_kwargs={"ombu_channels": 0, "stochastic_gates": False},
    )
    assert "val_asymmetric" in fitted.history
    assert np.isfinite(fitted.history["val_asymmetric"]).all()
python -m pytest packages/omnibias-torch/tests/test_joint_operator.py -k test_joint_operator_fit_accepts_asymmetric_loss -q
View source ↗
predictmethod
predict(self, x: NDArray[Any]) -> NDArray[np.float32]

See the class contract above and the exact method implementation below.

Parameters

NameType & defaultContract
xpositional-or-keywordNDArray[Any]RequiredSee the operation contract and implementation below.

Returns: NDArray[np.float32]

Executed regression · joint operator fit accepts sample weights

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_accepts_sample_weights() -> None:
    x, y = _synthetic(n=500)
    weights = np.ones(350, dtype=np.float32)
    weights[y[:350] < np.median(y[:350])] = 2.0
    fitted = fit_joint_operator_regressor(
        x[:350],
        y[:350],
        x[350:],
        y[350:],
        seed=1,
        epochs=80,
        patience=30,
        train_sample_weight=weights,
        val_sample_weight=np.ones(150, dtype=np.float32),
        model_kwargs={"ombu_channels": 0, "stochastic_gates": False},
    )
    pred = fitted.predict(x[350:])
    assert np.isfinite(pred).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_fit_accepts_sample_weights" -q
Method implementation
    def predict(self, x: NDArray[Any]) -> NDArray[np.float32]:
        self.model.eval()
        # Place the input on the model's device so predict works after the
        # model has been moved to a GPU (CPU-safe no-op otherwise). The
        # ``torch.empty(0)`` default keeps a parameterless model on CPU.
        device = next(self.model.parameters(), torch.empty(0)).device
        xs = (np.asarray(x, dtype=np.float32) - self.x_mean) / self.x_scale
        with torch.no_grad():
            xt = torch.as_tensor(xs, dtype=torch.float32, device=device)
            pred_z: NDArray[np.float32] = self.model(xt).cpu().numpy()
        return pred_z * self.y_scale + self.y_mean
View source ↗
selected_operatorsmethod
selected_operators(self, *, threshold: float=0.2, top_k: int | None=None) -> list[dict[str, Any]]

See the class contract above and the exact method implementation below.

Parameters

NameType & defaultContract
thresholdkeyword-onlyfloatDefault: 0.2See the operation contract and implementation below.
top_kkeyword-onlyint | NoneDefault: NoneSee the operation contract and implementation below.

Returns: list[dict[str, Any]]

Executed regression · joint operator fit recovers synthetic law primitives

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_recovers_synthetic_law_primitives() -> None:
    x, y = _synthetic()
    x_train, y_train = x[:500], y[:500]
    x_val, y_val = x[500:], y[500:]
    fitted = fit_joint_operator_regressor(
        x_train,
        y_train,
        x_val,
        y_val,
        seed=0,
        epochs=260,
        patience=80,
        sparsity_weight=2e-3,
        model_kwargs={
            "ombu_channels": 0,
            "stochastic_gates": False,
            "initial_gate_logit": 0.0,
        },
    )
    pred = fitted.predict(x_val)
    rmse = float(np.sqrt(np.mean((pred - y_val) ** 2)))
    selected_names = {row["name"] for row in fitted.selected_operators(top_k=12)}
    assert rmse < 0.25
    assert "x1^2" in selected_names
    assert "x2*x3" in selected_names
    assert "sin(x4)" in selected_names
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_fit_recovers_synthetic_law_primitives" -q
Method implementation
    def selected_operators(self, *, threshold: float = 0.2, top_k: int | None = None) -> list[dict[str, Any]]:
        return self.model.selected_operators(threshold=threshold, top_k=top_k)
View source ↗
Read the implementation
class FittedJointOperatorRegressor:
    """A fitted joint operator model plus normalization state."""

    model: JointOperatorRegressor
    x_mean: NDArray[np.float32]
    x_scale: NDArray[np.float32]
    y_mean: float
    y_scale: float
    history: dict[str, list[float]]

    def predict(self, x: NDArray[Any]) -> NDArray[np.float32]:
        self.model.eval()
        # Place the input on the model's device so predict works after the
        # model has been moved to a GPU (CPU-safe no-op otherwise). The
        # ``torch.empty(0)`` default keeps a parameterless model on CPU.
        device = next(self.model.parameters(), torch.empty(0)).device
        xs = (np.asarray(x, dtype=np.float32) - self.x_mean) / self.x_scale
        with torch.no_grad():
            xt = torch.as_tensor(xs, dtype=torch.float32, device=device)
            pred_z: NDArray[np.float32] = self.model(xt).cpu().numpy()
        return pred_z * self.y_scale + self.y_mean

    def selected_operators(self, *, threshold: float = 0.2, top_k: int | None = None) -> list[dict[str, Any]]:
        return self.model.selected_operators(threshold=threshold, top_k=top_k)
Other public import paths
  • omnibias.torch.architectures.FittedJointOperatorRegressor
  • omnibias.torch.architectures.joint_operator.FittedJointOperatorRegressor
CLASS

JointOperatorRegressor​

omnibias.torch.architectures.JointOperatorRegressorSource ↗
class JointOperatorRegressor(nn.Module)

Differentiable operator bank with trainable selection gates.

The model intentionally starts simple: named deterministic operators provide interpretability, optional trainable OMBU derivative channels add latent operator capacity, and a sparse linear head keeps selected operators exportable.

Parameters

NameType & defaultContract
in_featurespositional-or-keywordintRequiredSee the operation contract and implementation below.
include_rawkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
include_unarykeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
include_pairwisekeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
include_nestedkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
max_pairwisekeyword-onlyintDefault: 128See the operation contract and implementation below.
ombu_channelskeyword-onlyintDefault: 0See the operation contract and implementation below.
ombu_basekeyword-onlystr | ActivationSpec[Tensor]Default: 'tanh'See the operation contract and implementation below.
gate_temperaturekeyword-onlyfloatDefault: 1.0See the operation contract and implementation below.
stochastic_gateskeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
initial_gate_logitkeyword-onlyfloatDefault: -0.25See the operation contract and implementation below.
Executed regression · joint operator complexity loss is differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_complexity_loss_is_differentiable" -q
Regression example · joint operator bank contains scientific primitives

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_bank_contains_scientific_primitives() -> None:
    model = JointOperatorRegressor(4, ombu_channels=2, stochastic_gates=False)
    names = [item.name for item in model.operator_metadata]
    assert "x1^2" in names
    assert "x2*x3" in names
    assert "sin(x4)" in names
    assert "tanh(x1^2)" in names
    assert "ombu_grad_0" in names
    bank = model.operator_bank(torch.randn(5, 4))
    assert bank.shape == (5, model.n_operators)
python -m pytest packages/omnibias-torch/tests/test_joint_operator.py -k test_joint_operator_bank_contains_scientific_primitives -q

Fields & methods

__init__method
__init__(self, in_features: int, *, include_raw: bool=True, include_unary: bool=True, include_pairwise: bool=True, include_nested: bool=True, max_pairwise: int=128, ombu_channels: int=0, ombu_base: str | ActivationSpec[Tensor]='tanh', gate_temperature: float=1.0, stochastic_gates: bool=True, initial_gate_logit: float=-0.25) -> None

See the class contract above and the exact method implementation below.

Parameters

NameType & defaultContract
in_featurespositional-or-keywordintRequiredSee the operation contract and implementation below.
include_rawkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
include_unarykeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
include_pairwisekeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
include_nestedkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
max_pairwisekeyword-onlyintDefault: 128See the operation contract and implementation below.
ombu_channelskeyword-onlyintDefault: 0See the operation contract and implementation below.
ombu_basekeyword-onlystr | ActivationSpec[Tensor]Default: 'tanh'See the operation contract and implementation below.
gate_temperaturekeyword-onlyfloatDefault: 1.0See the operation contract and implementation below.
stochastic_gateskeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
initial_gate_logitkeyword-onlyfloatDefault: -0.25See the operation contract and implementation below.

Returns: None

Validation & failure cases
  • When in_features < 1, raises ValueError(f'in_features must be >= 1, got {in_features}').
  • When max_pairwise < 0, raises ValueError(f'max_pairwise must be >= 0, got {max_pairwise}').
  • When ombu_channels < 0, raises ValueError(f'ombu_channels must be >= 0, got {ombu_channels}').
  • When gate_temperature <= 0, raises ValueError(f'gate_temperature must be > 0, got {gate_temperature}').
  • When not self.operator_metadata, raises ValueError('operator bank is empty').
  • When self.ombu_spec.fastpath is None, raises TypeError('ombu_base must provide a derivative fastpath for joint derivative operators').
Executed regression · joint operator complexity loss is differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_complexity_loss_is_differentiable" -q
Method implementation
    def __init__(
        self,
        in_features: int,
        *,
        include_raw: bool = True,
        include_unary: bool = True,
        include_pairwise: bool = True,
        include_nested: bool = True,
        max_pairwise: int = 128,
        ombu_channels: int = 0,
        ombu_base: str | ActivationSpec[Tensor] = "tanh",
        gate_temperature: float = 1.0,
        stochastic_gates: bool = True,
        initial_gate_logit: float = -0.25,
    ) -> None:
        super().__init__()
        if in_features < 1:
            raise ValueError(f"in_features must be >= 1, got {in_features}")
        if max_pairwise < 0:
            raise ValueError(f"max_pairwise must be >= 0, got {max_pairwise}")
        if ombu_channels < 0:
            raise ValueError(f"ombu_channels must be >= 0, got {ombu_channels}")
        if gate_temperature <= 0:
            raise ValueError(f"gate_temperature must be > 0, got {gate_temperature}")

        self.in_features = in_features
        self.include_raw = include_raw
        self.include_unary = include_unary
        self.include_pairwise = include_pairwise
        self.include_nested = include_nested
        self.max_pairwise = max_pairwise
        self.ombu_channels = ombu_channels
        self.gate_temperature = gate_temperature
        self.stochastic_gates = stochastic_gates

        self.operator_metadata = _build_operator_metadata(
            in_features,
            include_raw=include_raw,
            include_unary=include_unary,
            include_pairwise=include_pairwise,
            include_nested=include_nested,
            max_pairwise=max_pairwise,
            ombu_channels=ombu_channels,
        )
        if not self.operator_metadata:
            raise ValueError("operator bank is empty")

        self.ombu_projection: nn.Linear | None
        self.ombu_spec: ActivationSpec[Tensor] | None
        if ombu_channels:
            self.ombu_projection = nn.Linear(in_features, ombu_channels)
            self.ombu_spec = ombu_base if isinstance(ombu_base, ActivationSpec) else get_activation(ombu_base)
            if self.ombu_spec.fastpath is None:
                raise TypeError("ombu_base must provide a derivative fastpath for joint derivative operators")
            try:
                self.ombu_spec.fastpath(torch.zeros(1), 2)
            except NotImplementedError as exc:
                raise TypeError("ombu_base must support derivative orders 1 and 2") from exc
        else:
            self.ombu_projection = None
            self.ombu_spec = None

        n_ops = len(self.operator_metadata)
        self.gate_logits = nn.Parameter(torch.full((n_ops,), float(initial_gate_logit)))
        self.readout = nn.Linear(n_ops, 1)
        complexity = torch.tensor([item.complexity for item in self.operator_metadata], dtype=torch.float32)
        self.operator_complexity: Tensor
        self.selection_readout_weight: Tensor
        self.register_buffer("operator_complexity", complexity)
        self.register_buffer("selection_readout_weight", torch.zeros(n_ops), persistent=False)
View source ↗
n_operatorsproperty
n_operators(self) -> int

See the class contract above and the exact method implementation below.

Returns: int

Executed regression · joint operator bank contains scientific primitives

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_bank_contains_scientific_primitives() -> None:
    model = JointOperatorRegressor(4, ombu_channels=2, stochastic_gates=False)
    names = [item.name for item in model.operator_metadata]
    assert "x1^2" in names
    assert "x2*x3" in names
    assert "sin(x4)" in names
    assert "tanh(x1^2)" in names
    assert "ombu_grad_0" in names
    bank = model.operator_bank(torch.randn(5, 4))
    assert bank.shape == (5, model.n_operators)
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_bank_contains_scientific_primitives" -q
Method implementation
    def n_operators(self) -> int:
        return len(self.operator_metadata)
View source ↗
gate_probabilitiesmethod
gate_probabilities(self) -> Tensor

See the class contract above and the exact method implementation below.

Returns: Tensor

Executed regression · joint operator complexity loss is differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_complexity_loss_is_differentiable" -q
Method implementation
    def gate_probabilities(self) -> Tensor:
        return torch.sigmoid(self.gate_logits / self.gate_temperature)
View source ↗
gatesmethod
gates(self) -> Tensor

See the class contract above and the exact method implementation below.

Returns: Tensor

Executed regression · joint operator complexity loss is differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_complexity_loss_is_differentiable" -q
Method implementation
    def gates(self) -> Tensor:
        if self.training and self.stochastic_gates:
            u = torch.rand_like(self.gate_logits).clamp_(1e-6, 1.0 - 1e-6)
            logistic_noise = torch.log(u) - torch.log1p(-u)
            return torch.sigmoid((self.gate_logits + logistic_noise) / self.gate_temperature)
        return self.gate_probabilities()
View source ↗
operator_bankmethod
operator_bank(self, x: Tensor) -> Tensor

See the class contract above and the exact method implementation below.

Parameters

NameType & defaultContract
xpositional-or-keywordTensorRequiredSee the operation contract and implementation below.

Returns: Tensor

Validation & failure cases
  • When x.ndim != 2 or x.shape[1] != self.in_features, raises ValueError(f'x must have shape (batch, {self.in_features}), got {tuple(x.shape)}').
Executed regression · joint operator complexity loss is differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_complexity_loss_is_differentiable" -q
Method implementation
    def operator_bank(self, x: Tensor) -> Tensor:
        if x.ndim != 2 or x.shape[1] != self.in_features:
            raise ValueError(f"x must have shape (batch, {self.in_features}), got {tuple(x.shape)}")

        cols: list[Tensor] = []
        if self.include_raw:
            cols.extend([x[:, j] for j in range(self.in_features)])

        if self.include_unary:
            for j in range(self.in_features):
                col = x[:, j]
                cols.extend(
                    [
                        col.square(),
                        torch.tanh(col),
                        torch.sin(col),
                        torch.cos(col),
                        torch.log1p(col.abs()),
                        1.0 / (1.0 + col.abs()),
                        torch.exp(col.clamp(-3.0, 3.0)),
                    ]
                )

        if self.include_pairwise and self.max_pairwise:
            for idx, (j, k) in enumerate(combinations(range(self.in_features), 2)):
                if idx >= self.max_pairwise:
                    break
                cols.append(x[:, j] * x[:, k])

        if self.include_nested:
            for j in range(self.in_features):
                col = x[:, j]
                square = col.square()
                cols.extend([torch.tanh(square), torch.sin(square), torch.log1p(square)])

        if self.ombu_projection is not None and self.ombu_spec is not None:
            z = self.ombu_projection(x)
            fp = self.ombu_spec.fastpath
            assert fp is not None
            cols.extend([self.ombu_spec.forward(z[:, h]) for h in range(self.ombu_channels)])
            cols.extend([fp(z[:, h], 1) for h in range(self.ombu_channels)])
            cols.extend([fp(z[:, h], 2) for h in range(self.ombu_channels)])

        return torch.stack(cols, dim=1)
View source ↗
forwardmethod
forward(self, x: Tensor) -> Tensor

See the class contract above and the exact method implementation below.

Parameters

NameType & defaultContract
xpositional-or-keywordTensorRequiredSee the operation contract and implementation below.

Returns: Tensor

Executed regression · joint operator complexity loss is differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_complexity_loss_is_differentiable" -q
Method implementation
    def forward(self, x: Tensor) -> Tensor:
        bank = self.operator_bank(x)
        gated = bank * self.gates()
        out: Tensor = self.readout(gated).squeeze(-1)
        return out
View source ↗
complexity_lossmethod
complexity_loss(self) -> Tensor

See the class contract above and the exact method implementation below.

Returns: Tensor

Executed regression · joint operator complexity loss is differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_complexity_loss_is_differentiable" -q
Method implementation
    def complexity_loss(self) -> Tensor:
        probs = self.gate_probabilities()
        return (probs * self.operator_complexity).mean()
View source ↗
active_operator_countmethod
active_operator_count(self, threshold: float=0.2) -> int

See the class contract above and the exact method implementation below.

Parameters

NameType & defaultContract
thresholdpositional-or-keywordfloatDefault: 0.2See the operation contract and implementation below.

Returns: int

Method implementation
    def active_operator_count(self, threshold: float = 0.2) -> int:
        return int((self.gate_probabilities().detach() >= threshold).sum().item())
View source ↗
selected_operatorsmethod
selected_operators(self, *, threshold: float=0.2, top_k: int | None=None) -> list[dict[str, Any]]

See the class contract above and the exact method implementation below.

Parameters

NameType & defaultContract
thresholdkeyword-onlyfloatDefault: 0.2See the operation contract and implementation below.
top_kkeyword-onlyint | NoneDefault: NoneSee the operation contract and implementation below.

Returns: list[dict[str, Any]]

Executed regression · joint operator fit recovers synthetic law primitives

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_recovers_synthetic_law_primitives() -> None:
    x, y = _synthetic()
    x_train, y_train = x[:500], y[:500]
    x_val, y_val = x[500:], y[500:]
    fitted = fit_joint_operator_regressor(
        x_train,
        y_train,
        x_val,
        y_val,
        seed=0,
        epochs=260,
        patience=80,
        sparsity_weight=2e-3,
        model_kwargs={
            "ombu_channels": 0,
            "stochastic_gates": False,
            "initial_gate_logit": 0.0,
        },
    )
    pred = fitted.predict(x_val)
    rmse = float(np.sqrt(np.mean((pred - y_val) ** 2)))
    selected_names = {row["name"] for row in fitted.selected_operators(top_k=12)}
    assert rmse < 0.25
    assert "x1^2" in selected_names
    assert "x2*x3" in selected_names
    assert "sin(x4)" in selected_names
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_fit_recovers_synthetic_law_primitives" -q
Method implementation
    def selected_operators(self, *, threshold: float = 0.2, top_k: int | None = None) -> list[dict[str, Any]]:
        with torch.no_grad():
            probs = self.gate_probabilities().detach().cpu()
            weights = self.readout.weight.detach().squeeze(0).cpu()
            selection_weights = self.selection_readout_weight.detach().cpu()
            if torch.count_nonzero(selection_weights) == 0:
                selection_weights = weights
            contribution = (probs * weights.abs()).numpy()
            selection_score = (
                probs
                * selection_weights.abs()
                * torch.tensor([item.interpretability for item in self.operator_metadata])
                / torch.sqrt(self.operator_complexity.detach().cpu())
            ).numpy()
        rows = []
        for idx, meta in enumerate(self.operator_metadata):
            if probs[idx].item() >= threshold:
                rows.append(
                    _selected_payload(
                        idx,
                        meta,
                        probs[idx].item(),
                        weights[idx].item(),
                        contribution[idx],
                        selection_score[idx],
                        selection_weights[idx].item(),
                    )
                )
        if top_k is not None:
            ranked = np.argsort(-selection_score)[:top_k]
            by_idx = {row["index"]: row for row in rows}
            for idx in ranked:
                if int(idx) not in by_idx:
                    meta = self.operator_metadata[int(idx)]
                    rows.append(
                        _selected_payload(
                            int(idx),
                            meta,
                            probs[int(idx)].item(),
                            weights[int(idx)].item(),
                            contribution[int(idx)],
                            selection_score[int(idx)],
                            selection_weights[int(idx)].item(),
                        )
                    )
            rows = sorted(rows, key=lambda row: row["selection_score"], reverse=True)[:top_k]
        else:
            rows = sorted(rows, key=lambda row: row["selection_score"], reverse=True)
        return rows
View source ↗
in_featuresattribute
in_features: int = in_features

Public instance attribute assigned by __init__. The source assignment is in_features = in_features; it describes instance state, not a literal constructor default.

View source ↗
include_rawattribute
include_raw: bool = include_raw

Public instance attribute assigned by __init__. The source assignment is include_raw = include_raw; it describes instance state, not a literal constructor default.

View source ↗
include_unaryattribute
include_unary: bool = include_unary

Public instance attribute assigned by __init__. The source assignment is include_unary = include_unary; it describes instance state, not a literal constructor default.

View source ↗
include_pairwiseattribute
include_pairwise: bool = include_pairwise

Public instance attribute assigned by __init__. The source assignment is include_pairwise = include_pairwise; it describes instance state, not a literal constructor default.

View source ↗
include_nestedattribute
include_nested: bool = include_nested

Public instance attribute assigned by __init__. The source assignment is include_nested = include_nested; it describes instance state, not a literal constructor default.

View source ↗
max_pairwiseattribute
max_pairwise: int = max_pairwise

Public instance attribute assigned by __init__. The source assignment is max_pairwise = max_pairwise; it describes instance state, not a literal constructor default.

View source ↗
ombu_channelsattribute
ombu_channels: int = ombu_channels

Public instance attribute assigned by __init__. The source assignment is ombu_channels = ombu_channels; it describes instance state, not a literal constructor default.

View source ↗
gate_temperatureattribute
gate_temperature: float = gate_temperature

Public instance attribute assigned by __init__. The source assignment is gate_temperature = gate_temperature; it describes instance state, not a literal constructor default.

View source ↗
stochastic_gatesattribute
stochastic_gates: bool = stochastic_gates

Public instance attribute assigned by __init__. The source assignment is stochastic_gates = stochastic_gates; it describes instance state, not a literal constructor default.

View source ↗
operator_metadataattribute
operator_metadata = _build_operator_metadata(in_features, include_raw=include_raw, include_unary=include_unary, include_pairwise=include_pairwise, include_nested=include_nested, max_pairwise=max_pairwise, ombu_channels=ombu_channels)

Public instance attribute assigned by __init__. The source assignment is operator_metadata = _build_operator_metadata(in_features, include_raw=include_raw, include_unary=include_unary, include_pairwise=include_pairwise, include_nested=include_nested, max_pairwise=max_pairwise, ombu_channels=ombu_channels); it describes instance state, not a literal constructor default.

Regression example · joint operator bank contains scientific primitives

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_bank_contains_scientific_primitives() -> None:
    model = JointOperatorRegressor(4, ombu_channels=2, stochastic_gates=False)
    names = [item.name for item in model.operator_metadata]
    assert "x1^2" in names
    assert "x2*x3" in names
    assert "sin(x4)" in names
    assert "tanh(x1^2)" in names
    assert "ombu_grad_0" in names
    bank = model.operator_bank(torch.randn(5, 4))
    assert bank.shape == (5, model.n_operators)
python -m pytest packages/omnibias-torch/tests/test_joint_operator.py -k test_joint_operator_bank_contains_scientific_primitives -q
View source ↗
ombu_projectionattribute
ombu_projection: nn.Linear | None

Public instance attribute assigned by __init__. The source assignment is ombu_projection = null; it describes instance state, not a literal constructor default.

View source ↗
ombu_specattribute
ombu_spec: ActivationSpec[Tensor] | None

Public instance attribute assigned by __init__. The source assignment is ombu_spec = null; it describes instance state, not a literal constructor default.

View source ↗
gate_logitsattribute
gate_logits = nn.Parameter(torch.full((n_ops,), float(initial_gate_logit)))

Public instance attribute assigned by __init__. The source assignment is gate_logits = nn.Parameter(torch.full((n_ops,), float(initial_gate_logit))); it describes instance state, not a literal constructor default.

Regression example · joint operator complexity loss is differentiable

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_complexity_loss_is_differentiable() -> None:
    model = JointOperatorRegressor(3, ombu_channels=1, stochastic_gates=True)
    x = torch.randn(16, 3)
    y = torch.randn(16)
    pred = model(x)
    loss = torch.mean((pred - y) ** 2) + 1e-3 * model.complexity_loss()
    loss.backward()
    assert model.gate_logits.grad is not None
    assert torch.isfinite(model.gate_logits.grad).all()
python -m pytest packages/omnibias-torch/tests/test_joint_operator.py -k test_joint_operator_complexity_loss_is_differentiable -q
View source ↗
readoutattribute
readout = nn.Linear(n_ops, 1)

Public instance attribute assigned by __init__. The source assignment is readout = nn.Linear(n_ops, 1); it describes instance state, not a literal constructor default.

View source ↗
operator_complexityattribute
operator_complexity: Tensor

Public instance attribute assigned by __init__. The source assignment is operator_complexity = null; it describes instance state, not a literal constructor default.

View source ↗
selection_readout_weightattribute
selection_readout_weight: Tensor

Public instance attribute assigned by __init__. The source assignment is selection_readout_weight = null; it describes instance state, not a literal constructor default.

View source ↗
Read the implementation
class JointOperatorRegressor(nn.Module):
    """Differentiable operator bank with trainable selection gates.

    The model intentionally starts simple: named deterministic operators
    provide interpretability, optional trainable OMBU derivative channels add
    latent operator capacity, and a sparse linear head keeps selected operators
    exportable.
    """

    def __init__(
        self,
        in_features: int,
        *,
        include_raw: bool = True,
        include_unary: bool = True,
        include_pairwise: bool = True,
        include_nested: bool = True,
        max_pairwise: int = 128,
        ombu_channels: int = 0,
        ombu_base: str | ActivationSpec[Tensor] = "tanh",
        gate_temperature: float = 1.0,
        stochastic_gates: bool = True,
        initial_gate_logit: float = -0.25,
    ) -> None:
        super().__init__()
        if in_features < 1:
            raise ValueError(f"in_features must be >= 1, got {in_features}")
        if max_pairwise < 0:
            raise ValueError(f"max_pairwise must be >= 0, got {max_pairwise}")
        if ombu_channels < 0:
            raise ValueError(f"ombu_channels must be >= 0, got {ombu_channels}")
        if gate_temperature <= 0:
            raise ValueError(f"gate_temperature must be > 0, got {gate_temperature}")

        self.in_features = in_features
        self.include_raw = include_raw
        self.include_unary = include_unary
        self.include_pairwise = include_pairwise
        self.include_nested = include_nested
        self.max_pairwise = max_pairwise
        self.ombu_channels = ombu_channels
        self.gate_temperature = gate_temperature
        self.stochastic_gates = stochastic_gates

        self.operator_metadata = _build_operator_metadata(
            in_features,
            include_raw=include_raw,
            include_unary=include_unary,
            include_pairwise=include_pairwise,
            include_nested=include_nested,
            max_pairwise=max_pairwise,
            ombu_channels=ombu_channels,
        )
        if not self.operator_metadata:
            raise ValueError("operator bank is empty")

        self.ombu_projection: nn.Linear | None
        self.ombu_spec: ActivationSpec[Tensor] | None
        if ombu_channels:
            self.ombu_projection = nn.Linear(in_features, ombu_channels)
            self.ombu_spec = ombu_base if isinstance(ombu_base, ActivationSpec) else get_activation(ombu_base)
            if self.ombu_spec.fastpath is None:
                raise TypeError("ombu_base must provide a derivative fastpath for joint derivative operators")
            try:
                self.ombu_spec.fastpath(torch.zeros(1), 2)
            except NotImplementedError as exc:
                raise TypeError("ombu_base must support derivative orders 1 and 2") from exc
        else:
            self.ombu_projection = None
            self.ombu_spec = None

        n_ops = len(self.operator_metadata)
        self.gate_logits = nn.Parameter(torch.full((n_ops,), float(initial_gate_logit)))
        self.readout = nn.Linear(n_ops, 1)
        complexity = torch.tensor([item.complexity for item in self.operator_metadata], dtype=torch.float32)
        self.operator_complexity: Tensor
        self.selection_readout_weight: Tensor
        self.register_buffer("operator_complexity", complexity)
        self.register_buffer("selection_readout_weight", torch.zeros(n_ops), persistent=False)

    @property
    def n_operators(self) -> int:
        return len(self.operator_metadata)

    def gate_probabilities(self) -> Tensor:
        return torch.sigmoid(self.gate_logits / self.gate_temperature)

    def gates(self) -> Tensor:
        if self.training and self.stochastic_gates:
            u = torch.rand_like(self.gate_logits).clamp_(1e-6, 1.0 - 1e-6)
            logistic_noise = torch.log(u) - torch.log1p(-u)
            return torch.sigmoid((self.gate_logits + logistic_noise) / self.gate_temperature)
        return self.gate_probabilities()

    def operator_bank(self, x: Tensor) -> Tensor:
        if x.ndim != 2 or x.shape[1] != self.in_features:
            raise ValueError(f"x must have shape (batch, {self.in_features}), got {tuple(x.shape)}")

        cols: list[Tensor] = []
        if self.include_raw:
            cols.extend([x[:, j] for j in range(self.in_features)])

        if self.include_unary:
            for j in range(self.in_features):
                col = x[:, j]
                cols.extend(
                    [
                        col.square(),
                        torch.tanh(col),
                        torch.sin(col),
                        torch.cos(col),
                        torch.log1p(col.abs()),
                        1.0 / (1.0 + col.abs()),
                        torch.exp(col.clamp(-3.0, 3.0)),
                    ]
                )

        if self.include_pairwise and self.max_pairwise:
            for idx, (j, k) in enumerate(combinations(range(self.in_features), 2)):
                if idx >= self.max_pairwise:
                    break
                cols.append(x[:, j] * x[:, k])

        if self.include_nested:
            for j in range(self.in_features):
                col = x[:, j]
                square = col.square()
                cols.extend([torch.tanh(square), torch.sin(square), torch.log1p(square)])

        if self.ombu_projection is not None and self.ombu_spec is not None:
            z = self.ombu_projection(x)
            fp = self.ombu_spec.fastpath
            assert fp is not None
            cols.extend([self.ombu_spec.forward(z[:, h]) for h in range(self.ombu_channels)])
            cols.extend([fp(z[:, h], 1) for h in range(self.ombu_channels)])
            cols.extend([fp(z[:, h], 2) for h in range(self.ombu_channels)])

        return torch.stack(cols, dim=1)

    def forward(self, x: Tensor) -> Tensor:
        bank = self.operator_bank(x)
        gated = bank * self.gates()
        out: Tensor = self.readout(gated).squeeze(-1)
        return out

    def complexity_loss(self) -> Tensor:
        probs = self.gate_probabilities()
        return (probs * self.operator_complexity).mean()

    def active_operator_count(self, threshold: float = 0.2) -> int:
        return int((self.gate_probabilities().detach() >= threshold).sum().item())

    def selected_operators(self, *, threshold: float = 0.2, top_k: int | None = None) -> list[dict[str, Any]]:
        with torch.no_grad():
            probs = self.gate_probabilities().detach().cpu()
            weights = self.readout.weight.detach().squeeze(0).cpu()
            selection_weights = self.selection_readout_weight.detach().cpu()
            if torch.count_nonzero(selection_weights) == 0:
                selection_weights = weights
            contribution = (probs * weights.abs()).numpy()
            selection_score = (
                probs
                * selection_weights.abs()
                * torch.tensor([item.interpretability for item in self.operator_metadata])
                / torch.sqrt(self.operator_complexity.detach().cpu())
            ).numpy()
        rows = []
        for idx, meta in enumerate(self.operator_metadata):
            if probs[idx].item() >= threshold:
                rows.append(
                    _selected_payload(
                        idx,
                        meta,
                        probs[idx].item(),
                        weights[idx].item(),
                        contribution[idx],
                        selection_score[idx],
                        selection_weights[idx].item(),
                    )
                )
        if top_k is not None:
            ranked = np.argsort(-selection_score)[:top_k]
            by_idx = {row["index"]: row for row in rows}
            for idx in ranked:
                if int(idx) not in by_idx:
                    meta = self.operator_metadata[int(idx)]
                    rows.append(
                        _selected_payload(
                            int(idx),
                            meta,
                            probs[int(idx)].item(),
                            weights[int(idx)].item(),
                            contribution[int(idx)],
                            selection_score[int(idx)],
                            selection_weights[int(idx)].item(),
                        )
                    )
            rows = sorted(rows, key=lambda row: row["selection_score"], reverse=True)[:top_k]
        else:
            rows = sorted(rows, key=lambda row: row["selection_score"], reverse=True)
        return rows
Related APIsomnibias.torch.get_activation↗
Other public import paths
  • omnibias.torch.architectures.JointOperatorRegressor
  • omnibias.torch.architectures.joint_operator.JointOperatorRegressor
CLASS

OperatorMetadata​

omnibias.torch.architectures.OperatorMetadataSource ↗
class OperatorMetadata

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Description of one candidate operator in the joint bank.

Parameters

NameType & defaultContract
namedataclass fieldstrRequiredSee the operation contract and implementation below.
familydataclass fieldstrRequiredSee the operation contract and implementation below.
complexitydataclass fieldfloatRequiredSee the operation contract and implementation below.
inputsdataclass fieldtuple[int, ...]RequiredSee the operation contract and implementation below.
interpretabilitydataclass fieldfloatDefault: 1.0See the operation contract and implementation below.

Fields & methods

namefield
name: str

Declared str field on OperatorMetadata. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
familyfield
family: str

Declared str field on OperatorMetadata. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
complexityfield
complexity: float

Declared float field on OperatorMetadata. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
inputsfield
inputs: tuple[int, ...]

Declared tuple[int, ...] field on OperatorMetadata. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
interpretabilityfield
interpretability: float = 1.0

Declared float field on OperatorMetadata. Default expression: 1.0.

View source ↗
Read the implementation
class OperatorMetadata:
    """Description of one candidate operator in the joint bank."""

    name: str
    family: str
    complexity: float
    inputs: tuple[int, ...]
    interpretability: float = 1.0
Other public import paths
  • omnibias.torch.architectures.OperatorMetadata
  • omnibias.torch.architectures.joint_operator.OperatorMetadata
FUNCTION

fit_joint_operator_regressor​

omnibias.torch.architectures.fit_joint_operator_regressorSource ↗
fit_joint_operator_regressor(x_train: NDArray[Any], y_train: NDArray[Any], x_val: NDArray[Any] | None=None, y_val: NDArray[Any] | None=None, *, seed: int=0, epochs: int=400, lr: float=0.03, batch_size: int | None=None, patience: int=60, sparsity_weight: float=0.001, weight_decay: float=1e-05, train_sample_weight: NDArray[Any] | None=None, val_sample_weight: NDArray[Any] | None=None, asymmetric_weight: float=0.0, validation_asymmetric_weight: float | None=None, asymmetric_under_scale: float=13.0, asymmetric_over_scale: float=10.0, standardize_x: bool=False, polish_readout: bool=True, polish_ridge: float=1e-06, validation_selection_metric: Callable[[NDArray[Any], NDArray[Any]], float] | None=None, validation_selection_complexity_weight: float=0.0, model_kwargs: dict[str, Any] | None=None) -> FittedJointOperatorRegressor

Fit a JointOperatorRegressor with train/validation early stopping.

Parameters

NameType & defaultContract
x_trainpositional-or-keywordNDArray[Any]RequiredSee the operation contract and implementation below.
y_trainpositional-or-keywordNDArray[Any]RequiredSee the operation contract and implementation below.
x_valpositional-or-keywordNDArray[Any] | NoneDefault: NoneSee the operation contract and implementation below.
y_valpositional-or-keywordNDArray[Any] | NoneDefault: NoneSee the operation contract and implementation below.
seedkeyword-onlyintDefault: 0See the operation contract and implementation below.
epochskeyword-onlyintDefault: 400See the operation contract and implementation below.
lrkeyword-onlyfloatDefault: 0.03See the operation contract and implementation below.
batch_sizekeyword-onlyint | NoneDefault: NoneSee the operation contract and implementation below.
patiencekeyword-onlyintDefault: 60See the operation contract and implementation below.
sparsity_weightkeyword-onlyfloatDefault: 0.001See the operation contract and implementation below.
weight_decaykeyword-onlyfloatDefault: 1e-05See the operation contract and implementation below.
train_sample_weightkeyword-onlyNDArray[Any] | NoneDefault: NoneSee the operation contract and implementation below.
val_sample_weightkeyword-onlyNDArray[Any] | NoneDefault: NoneSee the operation contract and implementation below.
asymmetric_weightkeyword-onlyfloatDefault: 0.0See the operation contract and implementation below.
validation_asymmetric_weightkeyword-onlyfloat | NoneDefault: NoneSee the operation contract and implementation below.
asymmetric_under_scalekeyword-onlyfloatDefault: 13.0See the operation contract and implementation below.
asymmetric_over_scalekeyword-onlyfloatDefault: 10.0See the operation contract and implementation below.
standardize_xkeyword-onlyboolDefault: FalseSee the operation contract and implementation below.
polish_readoutkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
polish_ridgekeyword-onlyfloatDefault: 1e-06See the operation contract and implementation below.
validation_selection_metrickeyword-onlyCallable[[NDArray[Any], NDArray[Any]], float] | NoneDefault: NoneSee the operation contract and implementation below.
validation_selection_complexity_weightkeyword-onlyfloatDefault: 0.0See the operation contract and implementation below.
model_kwargskeyword-onlydict[str, Any] | NoneDefault: NoneSee the operation contract and implementation below.
ReturnsFittedJointOperatorRegressor
Explicit validation & failure cases
  • When x_train.ndim != 2, raises ValueError('x_train must be a 2D array').
  • When x_val.ndim != 2 or x_val.shape[1] != x_train.shape[1], raises ValueError('x_val must be a 2D array with the same feature count as x_train').
  • When asymmetric_weight < 0.0, raises ValueError(f'asymmetric_weight must be non-negative, got {asymmetric_weight}').
  • When validation_asymmetric_weight < 0.0, raises ValueError(f'validation_asymmetric_weight must be non-negative, got {validation_asymmetric_weight}').
  • When validation_selection_complexity_weight < 0.0, raises ValueError(f'validation_selection_complexity_weight must be non-negative, got {validation_selection_complexity_weight}').
  • When not np.isfinite(score), raises ValueError(f'validation_selection_metric returned a non-finite score: {score}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · joint operator fit accepts asymmetric loss

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_accepts_asymmetric_loss() -> None:
    x, y = _synthetic(n=420)
    fitted = fit_joint_operator_regressor(
        x[:300],
        y[:300],
        x[300:],
        y[300:],
        seed=2,
        epochs=60,
        patience=25,
        asymmetric_weight=1e-3,
        model_kwargs={"ombu_channels": 0, "stochastic_gates": False},
    )
    assert "val_asymmetric" in fitted.history
    assert np.isfinite(fitted.history["val_asymmetric"]).all()
python -m pytest "packages/omnibias-torch/tests/test_joint_operator.py::test_joint_operator_fit_accepts_asymmetric_loss" -q
Regression example · joint operator fit accepts sample weights

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗
def test_joint_operator_fit_accepts_sample_weights() -> None:
    x, y = _synthetic(n=500)
    weights = np.ones(350, dtype=np.float32)
    weights[y[:350] < np.median(y[:350])] = 2.0
    fitted = fit_joint_operator_regressor(
        x[:350],
        y[:350],
        x[350:],
        y[350:],
        seed=1,
        epochs=80,
        patience=30,
        train_sample_weight=weights,
        val_sample_weight=np.ones(150, dtype=np.float32),
        model_kwargs={"ombu_channels": 0, "stochastic_gates": False},
    )
    pred = fitted.predict(x[350:])
    assert np.isfinite(pred).all()
python -m pytest packages/omnibias-torch/tests/test_joint_operator.py -k test_joint_operator_fit_accepts_sample_weights -q
Read the implementation
def fit_joint_operator_regressor(
    x_train: NDArray[Any],
    y_train: NDArray[Any],
    x_val: NDArray[Any] | None = None,
    y_val: NDArray[Any] | None = None,
    *,
    seed: int = 0,
    epochs: int = 400,
    lr: float = 3e-2,
    batch_size: int | None = None,
    patience: int = 60,
    sparsity_weight: float = 1e-3,
    weight_decay: float = 1e-5,
    train_sample_weight: NDArray[Any] | None = None,
    val_sample_weight: NDArray[Any] | None = None,
    asymmetric_weight: float = 0.0,
    validation_asymmetric_weight: float | None = None,
    asymmetric_under_scale: float = 13.0,
    asymmetric_over_scale: float = 10.0,
    standardize_x: bool = False,
    polish_readout: bool = True,
    polish_ridge: float = 1e-6,
    validation_selection_metric: Callable[[NDArray[Any], NDArray[Any]], float] | None = None,
    validation_selection_complexity_weight: float = 0.0,
    model_kwargs: dict[str, Any] | None = None,
) -> FittedJointOperatorRegressor:
    """Fit a :class:`JointOperatorRegressor` with train/validation early stopping."""

    if x_val is None:
        x_val = x_train
    if y_val is None:
        y_val = y_train
    x_train = np.asarray(x_train, dtype=np.float32)
    y_train = np.asarray(y_train, dtype=np.float32)
    x_val = np.asarray(x_val, dtype=np.float32)
    y_val = np.asarray(y_val, dtype=np.float32)
    if x_train.ndim != 2:
        raise ValueError("x_train must be a 2D array")
    if x_val.ndim != 2 or x_val.shape[1] != x_train.shape[1]:
        raise ValueError("x_val must be a 2D array with the same feature count as x_train")
    if asymmetric_weight < 0.0:
        raise ValueError(f"asymmetric_weight must be non-negative, got {asymmetric_weight}")
    if validation_asymmetric_weight is None:
        validation_asymmetric_weight = asymmetric_weight
    if validation_asymmetric_weight < 0.0:
        raise ValueError(f"validation_asymmetric_weight must be non-negative, got {validation_asymmetric_weight}")
    if validation_selection_complexity_weight < 0.0:
        raise ValueError(
            "validation_selection_complexity_weight must be non-negative, "
            f"got {validation_selection_complexity_weight}"
        )
    train_weight = _prepare_sample_weight(train_sample_weight, x_train.shape[0], "train_sample_weight")
    val_weight = _prepare_sample_weight(val_sample_weight, x_val.shape[0], "val_sample_weight")

    if standardize_x:
        x_mean = x_train.mean(axis=0)
        x_scale = x_train.std(axis=0)
        x_scale = np.where(x_scale < 1e-6, 1.0, x_scale).astype(np.float32)
    else:
        x_mean = np.zeros(x_train.shape[1], dtype=np.float32)
        x_scale = np.ones(x_train.shape[1], dtype=np.float32)
    y_mean = float(y_train.mean())
    y_scale = float(y_train.std())
    if y_scale < 1e-6:
        y_scale = 1.0

    xs_train = (x_train - x_mean) / x_scale
    xs_val = (x_val - x_mean) / x_scale
    ys_train = (y_train - y_mean) / y_scale
    ys_val = (y_val - y_mean) / y_scale

    torch.manual_seed(seed)
    kwargs = model_kwargs or {}
    model = JointOperatorRegressor(x_train.shape[1], **kwargs)
    # The numpy I/O contract is float32 end-to-end, but the model's parameters
    # are created at the global default dtype. Pin the model to float32 so a
    # process-wide torch.set_default_dtype(torch.float64) cannot desync the
    # float64 params from the float32 inputs (which otherwise crashes the
    # ridge polish and any nn.Linear inside the operator bank).
    model.to(torch.float32)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    xtr = torch.as_tensor(xs_train, dtype=torch.float32)
    ytr = torch.as_tensor(ys_train, dtype=torch.float32)
    xv = torch.as_tensor(xs_val, dtype=torch.float32)
    yv = torch.as_tensor(ys_val, dtype=torch.float32)
    wtr = torch.as_tensor(train_weight, dtype=torch.float32)
    wv = torch.as_tensor(val_weight, dtype=torch.float32)

    history: dict[str, list[float]] = {
        "train_loss": [],
        "val_rmse_z": [],
        "val_asymmetric": [],
        "complexity": [],
        "val_selection_score": [],
    }
    best_score = float("inf")
    best_state: dict[str, Tensor] | None = None
    stale = 0
    n = xtr.shape[0]
    full_batch = batch_size is None or batch_size >= n

    for _ in range(epochs):
        model.train()
        if full_batch:
            batch_indices = [torch.arange(n)]
        else:
            assert batch_size is not None
            perm = torch.randperm(n)
            batch_indices = [perm[start : start + batch_size] for start in range(0, n, batch_size)]
        epoch_losses = []
        for idx in batch_indices:
            opt.zero_grad()
            pred = model(xtr[idx])
            mse = _weighted_mse(pred, ytr[idx], wtr[idx])
            asymmetric = _weighted_asymmetric_rul_loss(
                pred,
                ytr[idx],
                wtr[idx],
                y_scale=y_scale,
                under_scale=asymmetric_under_scale,
                over_scale=asymmetric_over_scale,
            )
            complexity = model.complexity_loss()
            loss = mse + asymmetric_weight * asymmetric + sparsity_weight * complexity
            loss.backward()  # type: ignore[no-untyped-call]
            opt.step()
            epoch_losses.append(float(loss.detach().cpu()))

        model.eval()
        with torch.no_grad():
            val_pred = model(xv)
            val_mse = _weighted_mse(val_pred, yv, wv)
            val_asymmetric = _weighted_asymmetric_rul_loss(
                val_pred,
                yv,
                wv,
                y_scale=y_scale,
                under_scale=asymmetric_under_scale,
                over_scale=asymmetric_over_scale,
            )
            val_rmse = float(torch.sqrt(val_mse).cpu())
            val_asymmetric_value = float(val_asymmetric.cpu())
            complexity_value = float(model.complexity_loss().cpu())
            val_pred_cycles = (val_pred.cpu().numpy() * y_scale + y_mean).astype(np.float64)
        history["train_loss"].append(float(np.mean(epoch_losses)))
        history["val_rmse_z"].append(val_rmse)
        history["val_asymmetric"].append(val_asymmetric_value)
        history["complexity"].append(complexity_value)

        if validation_selection_metric is None:
            score = (
                val_rmse
                + validation_asymmetric_weight * val_asymmetric_value
                + 0.05 * sparsity_weight * complexity_value
            )
        else:
            score = float(validation_selection_metric(y_val.astype(np.float64), val_pred_cycles))
            if not np.isfinite(score):
                raise ValueError(f"validation_selection_metric returned a non-finite score: {score}")
            score += validation_selection_complexity_weight * complexity_value
        history["val_selection_score"].append(float(score))
        if score < best_score - 1e-6:
            best_score = score
            best_state = {name: value.detach().cpu().clone() for name, value in model.state_dict().items()}
            stale = 0
        else:
            stale += 1
            if stale >= patience:
                break

    if best_state is not None:
        model.load_state_dict(best_state)
    with torch.no_grad():
        model.selection_readout_weight.copy_(model.readout.weight.detach().squeeze(0))
    if polish_readout:
        _polish_linear_readout(
            model,
            torch.cat([xtr, xv], dim=0),
            torch.cat([ytr, yv], dim=0),
            sample_weight=torch.cat([wtr, wv], dim=0),
            ridge=polish_ridge,
        )
    model.eval()
    return FittedJointOperatorRegressor(
        model=model,
        x_mean=x_mean.astype(np.float32),
        x_scale=x_scale.astype(np.float32),
        y_mean=y_mean,
        y_scale=y_scale,
        history=history,
    )
Related APIsomnibias.torch.architectures.FittedJointOperatorRegressor↗omnibias.torch.architectures.JointOperatorRegressor↗
Other public import paths
  • omnibias.torch.architectures.fit_joint_operator_regressor
  • omnibias.torch.architectures.joint_operator.fit_joint_operator_regressor
← Package overviewRunnable learning examples →