omnibias.partition.torch.weights
Trainable soft partitions of unity, regional experts, temperature hardening, and scoped membership bounds.
import omnibias.partition.torch.weightsModule contract & mathematical conventions
Torch soft partition-of-unity weights (algebraically corresponding twin of the numpy reference).
partition_weights_arrays is the autograd-friendly kernel (raw tensors in, weights
out) that a trainable bridge module differentiates through; partition_weights
evaluates a ~omnibias.partition._core.params.PartitionParams. Both reproduce
omnibias.partition._core.weights.partition_weights bit-for-bit (float64, parity
~1e-9).
Terminology: the split gate sigmoid(beta (W.x - t)) hardens as beta -> inf -- the
feasibility / temperature sense of "collapse", distinct from the **founding bias collapse**
(the multi-bias delta -> 0 limit to the closed-form derivative sigma^(K-1); see
docs/derivatives.md).
Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.
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.
combine
omnibias.partition.torch.combineSource ↗combine(weights: Tensor, region_outputs: Tensor) -> TensorBlend per-region outputs by partition weights: F = sum_l w_l out_l (torch twin).
Blend region predictions using already-computed partition weights. Scalar expert outputs produce one scalar per sample; vector expert outputs retain their final component axis. The blend is differentiable with respect to both tensors on supported backend paths. This function consumes the supplied weights; it does not independently construct or certify a partition of unity.
Parameters
| Name | Type & default | Contract |
|---|---|---|
weightspositional-or-keyword | TensorRequired | Region weights, shape (B, L), where L is the number of regions. |
region_outputspositional-or-keyword | TensorRequired | Scalar expert outputs (B, L), returning (B,), or vector expert outputs (B, L, K), returning (B, K). Batch and region axes must agree with weights. |
TensorComplete runnable example · soft regions
This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.
"""Train the split and its experts together. Requires omnibias-partition."""
import torch
from omnibias.partition.torch.weights import combine, partition_weights_arrays
x = torch.linspace(-1, 1, 12, dtype=torch.float64).reshape(-1, 1)
W = torch.tensor([[0.8]], dtype=torch.float64, requires_grad=True)
threshold = torch.tensor([0.1], dtype=torch.float64, requires_grad=True)
slopes = torch.tensor([0.5, -0.2], dtype=torch.float64, requires_grad=True)
weights = partition_weights_arrays(W, threshold, x, beta=3.0, depth=1)
experts = x * slopes[None, :]
prediction = combine(weights, experts)
assert weights.shape == (12, 2)
torch.testing.assert_close(weights.sum(-1), torch.ones(12, dtype=x.dtype))
assert bool((weights >= 0).all())
loss = (prediction - x[:, 0].square()).square().mean()
loss.backward()
for parameter in (W, threshold, slopes):
assert parameter.grad is not None and torch.isfinite(parameter.grad).all()
print("Partition of unity; gradients reach gates, thresholds, and experts.")python soft_regions.pyRead the implementation
def combine(weights: Tensor, region_outputs: Tensor) -> Tensor:
r"""Blend per-region outputs by partition weights: ``F = sum_l w_l out_l`` (torch twin)."""
if region_outputs.ndim == 2: # (n, L) scalar-per-region
return torch.einsum("nl,nl->n", weights, region_outputs)
return torch.einsum("nl,nlk->nk", weights, region_outputs)Other public import paths
omnibias.partition.torch.combineomnibias.partition.torch.weights.combine
partition_weights
omnibias.partition.torch.partition_weightsSource ↗partition_weights(params: PartitionParams, X: object, beta: float) -> TensorSoft partition weights (n, 2**depth) for a PartitionParams (float64).
Parameters
| Name | Type & default | Contract |
|---|---|---|
paramspositional-or-keyword | PartitionParamsRequired | See the operation contract and implementation below. |
Xpositional-or-keyword | objectRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatRequired | See the operation contract and implementation below. |
TensorExecuted regression · torch matches numpy
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.
def test_torch_matches_numpy(beta: float) -> None:
from omnibias.partition.torch import partition_weights as tw
params, X = _sample()
ref = partition_weights(params, X, beta)
got = tw(params, X, beta).detach().cpu().numpy()
assert np.max(np.abs(got - ref)) < 1e-9python -m pytest "packages/omnibias-partition/tests/test_parity.py::test_torch_matches_numpy[0.5]" -qRead the implementation
def partition_weights(params: PartitionParams, X: object, beta: float) -> Tensor:
r"""Soft partition weights ``(n, 2**depth)`` for a :class:`PartitionParams` (float64)."""
Xv = torch.as_tensor(np.asarray(X, dtype=np.float64), dtype=torch.float64)
W = torch.as_tensor(params.W, dtype=torch.float64)
t = torch.as_tensor(params.t, dtype=torch.float64)
return partition_weights_arrays(W, t, Xv, float(beta), params.depth)Other public import paths
omnibias.partition.torch.partition_weightsomnibias.partition.torch.weights.partition_weights
partition_weights_arrays
omnibias.partition.torch.partition_weights_arraysSource ↗partition_weights_arrays(W: Tensor, t: Tensor, X: Tensor, beta: float, depth: int) -> TensorRegion weights (n, 2**depth) from raw tensors -- the differentiable kernel.
Use the raw-array path when gate weights or thresholds must learn. It forms sigmoid(beta * (X W^T - t)), chooses the corresponding gate or its complement for each binary region code, and multiplies those factors across splits. There are 2**depth regions. Nonnegative weights sum to one up to floating-point error. Large inverse temperature can saturate gates and suppress gradients; validate the operating range of the model.
Parameters
| Name | Type & default | Contract |
|---|---|---|
Wpositional-or-keyword | TensorRequired | Gate coefficient matrix with shape (depth, D). Each row defines one shared binary split. |
tpositional-or-keyword | TensorRequired | Threshold vector with shape (depth,). |
Xpositional-or-keyword | TensorRequired | Input samples with shape (B, D). Gate tensors must use compatible dtype and device. |
betapositional-or-keyword | floatRequired | Finite inverse-temperature scale for the sigmoid gates. Increasing a positive beta sharpens the split; it changes gradient concentration. |
depthpositional-or-keyword | intRequired | Number of shared splits. It must match the leading dimension of W and the length of t; the output has shape (B, 2**depth). |
TensorComplete runnable example · soft regions
This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.
"""Train the split and its experts together. Requires omnibias-partition."""
import torch
from omnibias.partition.torch.weights import combine, partition_weights_arrays
x = torch.linspace(-1, 1, 12, dtype=torch.float64).reshape(-1, 1)
W = torch.tensor([[0.8]], dtype=torch.float64, requires_grad=True)
threshold = torch.tensor([0.1], dtype=torch.float64, requires_grad=True)
slopes = torch.tensor([0.5, -0.2], dtype=torch.float64, requires_grad=True)
weights = partition_weights_arrays(W, threshold, x, beta=3.0, depth=1)
experts = x * slopes[None, :]
prediction = combine(weights, experts)
assert weights.shape == (12, 2)
torch.testing.assert_close(weights.sum(-1), torch.ones(12, dtype=x.dtype))
assert bool((weights >= 0).all())
loss = (prediction - x[:, 0].square()).square().mean()
loss.backward()
for parameter in (W, threshold, slopes):
assert parameter.grad is not None and torch.isfinite(parameter.grad).all()
print("Partition of unity; gradients reach gates, thresholds, and experts.")python soft_regions.pyExecuted regression · torch matches numpy
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.
def test_torch_matches_numpy(beta: float) -> None:
from omnibias.partition.torch import partition_weights as tw
params, X = _sample()
ref = partition_weights(params, X, beta)
got = tw(params, X, beta).detach().cpu().numpy()
assert np.max(np.abs(got - ref)) < 1e-9python -m pytest "packages/omnibias-partition/tests/test_parity.py::test_torch_matches_numpy[0.5]" -qRead the implementation
def partition_weights_arrays(
W: Tensor, t: Tensor, X: Tensor, beta: float, depth: int
) -> Tensor:
r"""Region weights ``(n, 2**depth)`` from raw tensors -- the differentiable kernel."""
codes = torch.as_tensor(region_code_matrix(depth), dtype=X.dtype, device=X.device) # (L, D)
z = X @ W.transpose(-1, -2) - t[None, :] # (n, D)
g = torch.sigmoid(beta * z) # (n, D)
gexp = g[:, None, :] # (n, 1, D)
bexp = codes[None, :, :] # (1, L, D)
factors = bexp * gexp + (1.0 - bexp) * (1.0 - gexp) # (n, L, D)
return torch.prod(factors, dim=-1) # (n, L)Other public import paths
omnibias.partition.torch.partition_weights_arraysomnibias.partition.torch.weights.partition_weights_arrays