"""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.")