"""Mixed Taylor coefficients for a heat-equation residual at one point."""
import torch
from omnibias.torch.jet_mv import jet_partials, mlp_jet_mv

weights = torch.tensor([[0.5, 0.7], [-0.2, 0.3]], dtype=torch.float64,
                       requires_grad=True)
bias = torch.tensor([0.1, -0.1], dtype=torch.float64)
readout = torch.tensor([[0.6, -0.4]], dtype=torch.float64)
layers = [(weights, bias, "tanh"), (readout, None, None)]
point = torch.tensor([0.2, 0.4], dtype=torch.float64)  # coordinates (t, x)
jet = mlp_jet_mv(point, layers, order=2)
partials = jet_partials(jet, dim=2, order=2)
residual = partials[(1, 0)] - 0.1 * partials[(0, 2)]
residual.square().mean().backward()
assert partials[(1, 1)].shape == (1,)
assert weights.grad is not None and torch.isfinite(weights.grad).all()
print("Mixed partials:", sorted(partials), "finite parameter gradients: True")