跳转至

线性求解器与控制流

先定义一个把向量映射到向量的算子，再把它交给求解器。 从下面的三步 Poisson 示例开始；更复杂的动态邻域、迭代算法与非线性方程放在后面。

求解器的原生 Tensor 与 Torch 互操作

kernel 调用两种都接受——src/dst/edge 字段可以直接传 torch tensor，不需要包装。本页的高级原生 求解驱动器不同：linear_solve / nonlinear_solve 会把迭代循环捕获为 设备端控制流（gf_control.repeat / gf_control.while），因此向量 必须是 tg.Tensor。在边界处用 tg.from_torch(x) 把 torch 存储零 拷贝包进来；只有仍持有 Torch 存储的值才能用 .to_torch() 零拷贝返回。 原生求解结果需要 .to_torch(copy=True)，且要求连续存储。 共享存储本身不会把原生梯度图接入 PyTorch autograd。

用邻居求和求解 Poisson 方程

目标很简单：在一条线上求解 Poisson 方程， 两端固定为零，中间放 7 个未知节点。

\[ -u''(x)=1,\qquad u(0)=u(1)=0 \]

只需要定义“给定 u，如何算出 A(u)”。每个节点读取左右邻居， 自身取两倍，再除以网格间距；迭代循环交给求解器，不需要手写 CG。

下面是完整计算过程。linear_solve 来自仓库中的 examples/solvers.py，不是 tg 的内置 API。

import tiga as tg
from solvers import linear_solve  # Helper in examples/solvers.py, not tg API.


class Poisson(tg.MessagePassing):
    reducer = tg.sum()

    def edge(self, src, dst, edge):
        return -src.u

    def node(self, dst, incoming):
        return 2.0 * dst.u + incoming


# 1. Seven interior nodes; omitted boundary values are fixed to zero.
n = 7
h = 1.0 / (n + 1)
graph = tg.Graph.stencil((n,), ((-1,), (1,)))
kernel = Poisson()


def apply(u):
    return kernel(graph=graph, src={"u": u}, dst={"u": u}, edge={}) / h


# 2. Solve A(u) = b. The driver owns the iteration loop.
b = tg.tensor([h] * n, dtype=tg.float32)
u = linear_solve(apply, b, tolerance=1e-6, max_iterations=32)

# 3. Compare the nodal solution with x(1-x)/2.
exact = [0.5 * (i * h) * (1.0 - i * h) for i in range(1, n + 1)]
error = max(abs(a - e) for a, e in zip(u.tolist(), exact))

运行 python examples/fem_poisson_minimal.py，得到：

[0.0546875, 0.09375, 0.1171875, 0.125, 0.1171875, 0.09375, 0.0546875]
max error: 0.00e+00

这里的“无矩阵”指不构造刚度矩阵：Graph.stencil 给出邻接， edge() 取负的邻居值，node() 加上自身的两倍。 图仍有 CSR 拓扑，省去的是单独的刚度矩阵与逐边刚度数组。

本例使用 tg.Tensor，因为当前辅助求解器捕获的是原生控制流； 这不是普通消息传递必须更换 Torch Tensor 的要求。

为什么这是 FEM？边界怎么处理？

一维均匀网格上的分段线性 P1 有限元 给出以下离散方程。h 是节点间距，b 是单位载荷对应的右端项：

\[ (Au)_i=\frac{2u_i-u_{i-1}-u_{i+1}}{h},\qquad b_i=h \]

图只存内部节点，越界邻居被省略。被省略的边界值本来就是零， 但 node() 中的对角系数仍是 2，不能改成有效邻居数。

一维 P1 stencil：内部节点读取左右邻居，边界固定为零。

解析解用于检查节点结果：

\[ u_{\mathrm{exact}}(x)=\frac{x(1-x)}{2} \]
停止条件、固定迭代与求导

tolerance=1e-6 检查绝对残差范数，max_iterations=32 给出迭代上界； 达到上界不等于保证收敛，因此示例额外检查解析解误差。

固定次数迭代和对载荷求导保留在 进阶示例 fem_poisson.py。 控制流与梯度契约见线性求解器。

动态半径图上的无矩阵线性求解

是什么。 这里完全没有网格文件——图是从数据生成的。给定平面上 N 个点，距离不超过截断半径 r 的两个不同点成为邻居（半径图）。在该图上， 示例求解 (m·I + L) u = b，其中 L 是图拉普拉斯算子：对每个点，把它 自身的值与每个邻居的值之差求和。质量偏移 m = 1 保证算子正定， 右端项是周期模式 b_i = 1 + (i mod 3)。

半径关系：平面上的两个不同点距离不超过截断半径 r 时成为邻居；点 i 处的算子是 m·u_i 加上对所有邻居的 u_i − u_j 求和

\[ \big((mI + L)\,u\big)_i \;=\; m\,u_i \;+ \sum_{j\ne i \,:\, \lVert p_i - p_j \rVert \le r} (u_i - u_j) \;=\; b_i, \qquad b_i = 1 + (i \bmod 3) \]

示例把点放在间距为 h 的直线上，并设 r = 1.01·h，因此每个点的邻域 恰好是相邻的两个点，生成的关系是一条路径。

求解器和算子接口与 FEM 情形完全相同：拓扑由 Graph.radius 产生， 关系的实现与复用是 MessagePassing UDF 之下 provider 层的决策。 CG 在一个 gf_control.while 中捕获收敛判断和四个随迭代携带的状态； 用户不需要写任何邻接结构或稀疏矩阵。

class ShiftedRadiusLaplacian(tg.MessagePassing):
    """Apply ``mass * u + sum_neighbour(u_dst - u_src)``."""

    reducer = tg.sum()

    def edge(self, src, dst, edge):
        return dst.u - src.u

    def node(self, dst, laplacian, mass):
        return mass * dst.u + laplacian


def solve(
    points: int = 8,
    *,
    tolerance: float = 1.0e-5,
    max_iterations: int = 32,
):
    """Solve a positive-definite operator on a generated path relation."""
    if points < 2:
        raise ValueError("points must be at least two")
    spacing = 1.0 / (points - 1)
    positions = tg.tensor(
        [[index * spacing, 0.0] for index in range(points)],
        dtype=tg.float32,
    )  # (N, 2)
    # The logical relation remains procedural. Its current physical snapshot
    # may be built, streamed, cached, or rebuilt by the selected provider.
    graph = tg.Graph.radius(positions, cutoff=1.01 * spacing)
    # A MessagePassing kernel bound to a Graph is the linear operator; the
    # solver binds the iterated vector to the "u" field each iteration.
    apply_laplacian = ShiftedRadiusLaplacian()
    rhs = tg.tensor(
        [1.0 + float(index % 3) for index in range(points)],
        dtype=tg.float32,
    )  # (N,)
    solution = linear_solve(
        apply_laplacian,
        rhs,
        method="cg",
        graph=graph,
        field="u",
        params={"mass": 1.0},
        tolerance=tolerance,
        max_iterations=max_iterations,
    )  # (N,)
    residual_norm = vector_norm(
        apply_laplacian(
            graph=graph,
            src={"u": solution},
            dst={"u": solution},
            mass=1.0,
        )
        - rhs
    )
    return solution, residual_norm, graph, apply_laplacian
完整源码：examples/meshfree_linear_solve.py（可直接运行） 
"""Dynamic-radius matrix-free solve with a bounded device-side CG loop."""

from __future__ import annotations

import tiga as tg
from solvers import linear_solve, vector_norm


class ShiftedRadiusLaplacian(tg.MessagePassing):
    """Apply ``mass * u + sum_neighbour(u_dst - u_src)``."""

    reducer = tg.sum()

    def edge(self, src, dst, edge):
        return dst.u - src.u

    def node(self, dst, laplacian, mass):
        return mass * dst.u + laplacian


def solve(
    points: int = 8,
    *,
    tolerance: float = 1.0e-5,
    max_iterations: int = 32,
):
    """Solve a positive-definite operator on a generated path relation."""
    if points < 2:
        raise ValueError("points must be at least two")
    spacing = 1.0 / (points - 1)
    positions = tg.tensor(
        [[index * spacing, 0.0] for index in range(points)],
        dtype=tg.float32,
    )  # (N, 2)
    # The logical relation remains procedural. Its current physical snapshot
    # may be built, streamed, cached, or rebuilt by the selected provider.
    graph = tg.Graph.radius(positions, cutoff=1.01 * spacing)
    # A MessagePassing kernel bound to a Graph is the linear operator; the
    # solver binds the iterated vector to the "u" field each iteration.
    apply_laplacian = ShiftedRadiusLaplacian()
    rhs = tg.tensor(
        [1.0 + float(index % 3) for index in range(points)],
        dtype=tg.float32,
    )  # (N,)
    solution = linear_solve(
        apply_laplacian,
        rhs,
        method="cg",
        graph=graph,
        field="u",
        params={"mass": 1.0},
        tolerance=tolerance,
        max_iterations=max_iterations,
    )  # (N,)
    residual_norm = vector_norm(
        apply_laplacian(
            graph=graph,
            src={"u": solution},
            dst={"u": solution},
            mass=1.0,
        )
        - rhs
    )
    return solution, residual_norm, graph, apply_laplacian


# Inspect: solve()[2].explain(), solve()[0].mlir(),
#   solve()[0].generated_code("cpu_loop")

if __name__ == "__main__":
    solve()
本机编译产物（Ryzen 7 255 · RTX 5070 Ti 实测） 
### ShiftedRadiusLaplacian
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: materialize-fixed-radius-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=3

### ShiftedRadiusLaplacian
backend: cuda:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: materialize-fixed-radius-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=3

共享求解器语法糖：linear_solve

是什么。 上面两个示例都用定常迭代求解 A·u = b：从一个初始猜测出发， 反复应用 A 来缩小残差 r = b − A·u。提供三种方法，由一个 method= 参数选择：

method= 算法 算子契约
"cg"（默认） 共轭梯度 对称正定
"bicgstab" BiCGStab 非对称
"richardson" 阻尼定步长迭代 任意算子，最慢

Richardson 采用朴素的校正步 u ← u + ω·r。当 A 对称正定时，共轭梯度 做得更好：沿一系列 A-正交搜索方向前进，在迭代之间携带四个状态 ——解 u、残差 r、搜索方向 p，以及内积 ρ = ⟨r, z⟩，其中 z 是经过 可选预条件处理的残差。BiCGStab 以每步两次算子应用为代价放弃了 对称性要求，额外携带影子残差 r̂ = r₀、算子像 v = A·p 和 t = A·s， 以及标量 ρ、α、ω。

定常求解器循环：算子、右端项与停止契约进入只捕获一次的循环体；状态 u、r、p 与 ρ 在迭代之间传递

\[ u_{k+1} \;=\; u_k + \omega\,(b - A u_k) \qquad\text{(Richardson)} \]
\[ \alpha_k = \frac{\langle r_k, z_k \rangle}{p_k^{\top} A\, p_k}, \qquad u_{k+1} = u_k + \alpha_k p_k, \qquad r_{k+1} = r_k - \alpha_k A p_k, \qquad p_{k+1} = z_{k+1} + \frac{\langle r_{k+1}, z_{k+1} \rangle}{\langle r_k, z_k \rangle}\, p_k \quad\text{(CG)} \]
\[ \rho_k = \langle \hat r, r_{k-1} \rangle, \qquad p_k = r_{k-1} + \underbrace{\tfrac{\rho_k}{\rho_{k-1}} \tfrac{\alpha_{k-1}}{\omega_{k-1}}}_{\beta}\,(p_{k-1} - \omega_{k-1} v_{k-1}), \qquad \alpha_k = \frac{\rho_k}{\langle \hat r, A p_k \rangle} \]
\[ s = r_{k-1} - \alpha_k A p_k, \qquad \omega_k = \frac{\langle A s,\, s \rangle}{\langle A s,\, A s \rangle}, \qquad u_k = u_{k-1} + \alpha_k p_k + \omega_k s, \qquad r_k = s - \omega_k A s \quad\text{(BiCGStab)} \]

solvers.py 用 gf_control 原语组合这些方法，循环写成 @tg.jit 下 普通的 Python for/while。每个示例把算子构造包在一个局部 solve() 里；求解器本身就是一次 linear_solve(operator, rhs, method="cg", …) 调用——没有别的叫 solve 的东西。该模块刻意不属于核心包——是语法糖，不是原语。

examples/solvers.py（完整源码） 
"""Matrix-free stationary solvers composed from Tiga primitives."""

from __future__ import annotations

from collections.abc import Callable, Mapping

import tiga as tg

Matvec = Callable[[tg.Tensor], tg.Tensor]
Preconditioner = Callable[[tg.Tensor], tg.Tensor]


def dot(left: tg.Tensor, right: tg.Tensor) -> tg.Tensor:
    """Capture a rank-one Hermitian dot product as Tensor algebra."""
    if not isinstance(left, tg.Tensor) or not isinstance(right, tg.Tensor):
        raise TypeError("dot operands must be tiga.Tensor values")
    if left.ndim != 1 or right.ndim != 1 or left.shape != right.shape:
        raise ValueError("dot requires rank-one operands with identical shapes")
    if left.dtype is not right.dtype or left.device != right.device:
        raise ValueError("dot operands must have identical dtype and device")
    return (left.conj() * right).sum()


def vector_norm(value: tg.Tensor) -> tg.Tensor:
    """Capture the Euclidean norm of one real vector.

    Complex norm needs a real-valued projection dtype that the minimal Tensor
    frontend does not expose yet, so it fails instead of returning a complex
    square root with misleading semantics.
    """
    if not isinstance(value, tg.Tensor) or value.ndim != 1:
        raise TypeError("vector_norm expects one rank-one tiga.Tensor")
    if value.dtype.kind != "float":
        raise TypeError("vector_norm currently requires a real floating Tensor")
    return dot(value, value).sqrt()


def _bind_operator(
    operator,
    *,
    graph: tg.Graph | None,
    field: str | None,
    src: Mapping[str, tg.Tensor] | None,
    dst: Mapping[str, tg.Tensor] | None,
    edge: Mapping[str, tg.Tensor] | None,
    params: Mapping[str, object] | None,
) -> tuple[Matvec, int | None]:
    """Return ``(matvec, extent)`` for a kernel or plain callable operator."""
    if isinstance(operator, tg.Kernel):
        if graph is None or field is None:
            raise TypeError(
                "a MessagePassing operator requires graph= and field= so the "
                "solver can bind the iterated vector")
        if not isinstance(field, str) or not field:
            raise TypeError("field must be a non-empty field name string")
        schema = graph.schema
        if schema.num_src != schema.num_dst:
            raise ValueError(
                "solvers require a square relation: graph has "
                f"{schema.num_src} sources and {schema.num_dst} destinations")
        constant_src = dict(src) if src is not None else {}
        constant_dst = dict(dst) if dst is not None else {}
        if field in constant_src or field in constant_dst:
            raise ValueError(
                f"field {field!r} is the iterated unknown and must not also "
                "be bound as a constant")
        constant_edge = dict(edge) if edge is not None else {}
        bound_params = dict(params) if params is not None else {}

        def apply(vector: tg.Tensor) -> tg.Tensor:
            return operator(
                graph=graph,
                src={**constant_src, field: vector},
                dst={**constant_dst, field: vector},
                edge=constant_edge,
                **bound_params,
            )

        return apply, schema.num_dst
    bindings = {
        "graph": graph, "field": field, "src": src,
        "dst": dst, "edge": edge, "params": params,
    }
    given = [name for name, value in bindings.items() if value is not None]
    if given:
        names = ", ".join(given)
        raise TypeError(
            f"{names}= only apply to MessagePassing operators; a plain "
            "callable closes over its own constants")
    if not callable(operator):
        raise TypeError(
            "operator must be a tiga.MessagePassing kernel or a "
            "Tensor -> Tensor callable")
    return operator, None


def _checked_apply(
    matvec: Matvec,
    vector: tg.Tensor,
    extent: int | None,
) -> tg.Tensor:
    """Apply one matrix-free product and validate the contract."""
    if extent is not None and vector.shape != (extent,):
        raise ValueError(
            f"operator expected input shape {(extent,)}, got {vector.shape}")
    output = matvec(vector)
    if not isinstance(output, tg.Tensor):
        raise TypeError("operator must return a tiga.Tensor")
    if output.shape != vector.shape:
        raise ValueError(
            f"operator must be square: input shape {vector.shape} produced "
            f"output shape {output.shape}")
    if output.dtype is not vector.dtype or output.device != vector.device:
        raise ValueError(
            "operator input and output must have identical dtype and device")
    return output


@tg.jit
def _richardson_loop(matvec, omega, state, rhs, iterations):
    for _ in range(iterations):
        state = state + omega * (rhs - matvec(state))
    return state


@tg.jit
def _cg_fixed_loop(
    matvec, prepare, solution, residual, direction, product, iterations,
):
    for _ in range(iterations):
        applied = matvec(direction)
        alpha = product / dot(direction, applied)
        solution = solution + alpha * direction
        residual = residual - alpha * applied
        preconditioned = prepare(residual)
        updated = dot(residual, preconditioned)
        direction = preconditioned + (updated / product) * direction
        product = updated
    return solution


@tg.jit
def _cg_tolerance_loop(
    matvec, prepare, solution, residual, direction, product,
    threshold_squared, max_iterations,
):
    for _ in range(max_iterations):
        if dot(residual, residual) <= threshold_squared:
            break
        applied = matvec(direction)
        alpha = product / dot(direction, applied)
        solution = solution + alpha * direction
        residual = residual - alpha * applied
        preconditioned = prepare(residual)
        updated = dot(residual, preconditioned)
        direction = preconditioned + (updated / product) * direction
        product = updated
    return solution


@tg.jit
def _bicgstab_fixed_loop(
    matvec, solution, residual, shadow, direction, velocity,
    rho, alpha, omega, iterations,
):
    for _ in range(iterations):
        rho_new = dot(shadow, residual)
        beta = (rho_new / rho) * (alpha / omega)
        direction = residual + beta * (direction - omega * velocity)
        velocity = matvec(direction)
        alpha = rho_new / dot(shadow, velocity)
        intermediate = residual - alpha * velocity
        second = matvec(intermediate)
        omega = dot(second, intermediate) / dot(second, second)
        solution = solution + alpha * direction + omega * intermediate
        residual = intermediate - omega * second
        rho = rho_new
    return solution


@tg.jit
def _bicgstab_tolerance_loop(
    matvec, solution, residual, shadow, direction, velocity,
    rho, alpha, omega, threshold_squared, max_iterations,
):
    for _ in range(max_iterations):
        if dot(residual, residual) <= threshold_squared:
            break
        rho_new = dot(shadow, residual)
        beta = (rho_new / rho) * (alpha / omega)
        direction = residual + beta * (direction - omega * velocity)
        velocity = matvec(direction)
        alpha = rho_new / dot(shadow, velocity)
        intermediate = residual - alpha * velocity
        second = matvec(intermediate)
        omega = dot(second, intermediate) / dot(second, second)
        solution = solution + alpha * direction + omega * intermediate
        residual = intermediate - omega * second
        rho = rho_new
    return solution


def richardson(
    operator,
    rhs: tg.Tensor,
    *,
    graph: tg.Graph | None = None,
    field: str | None = None,
    src: Mapping[str, tg.Tensor] | None = None,
    dst: Mapping[str, tg.Tensor] | None = None,
    edge: Mapping[str, tg.Tensor] | None = None,
    params: Mapping[str, object] | None = None,
    iterations: int,
    relaxation: tg.Tensor | float = 1.0,
    initial: tg.Tensor | None = None,
) -> tg.Tensor:
    """Capture a bounded matrix-free Richardson iteration.

    The body is captured once inside ``gf_control.repeat`` and can contain
    MessagePassing. It is intentionally fixed-count; use ``cg`` with
    ``tolerance`` for a device-side convergence test.
    """
    matvec, extent = _bind_operator(
        operator, graph=graph, field=field, src=src, dst=dst,
        edge=edge, params=params)
    if not isinstance(rhs, tg.Tensor) or rhs.ndim != 1:
        raise TypeError("richardson rhs must be a rank-one tiga.Tensor")
    if extent is not None and rhs.shape != (extent,):
        raise ValueError(
            f"richardson rhs shape must be {(extent,)}, got {rhs.shape}")
    if not isinstance(iterations, int) or isinstance(iterations, bool):
        raise TypeError("richardson iterations must be an integer")
    if iterations < 0:
        raise ValueError("richardson iterations must be non-negative")
    state = tg.zeros_like(rhs) if initial is None else initial
    if not isinstance(state, tg.Tensor) or state.shape != rhs.shape:
        raise ValueError("richardson initial value must match rhs shape")
    if state.dtype is not rhs.dtype or state.device != rhs.device:
        raise ValueError("richardson initial value must match rhs dtype and device")
    omega = relaxation
    return _richardson_loop(
        lambda value: _checked_apply(matvec, value, extent),
        omega,
        state,
        rhs,
        iterations,
    )


def cg(
    operator,
    rhs: tg.Tensor,
    *,
    graph: tg.Graph | None = None,
    field: str | None = None,
    src: Mapping[str, tg.Tensor] | None = None,
    dst: Mapping[str, tg.Tensor] | None = None,
    edge: Mapping[str, tg.Tensor] | None = None,
    params: Mapping[str, object] | None = None,
    iterations: int | None = None,
    tolerance: tg.Tensor | float | None = None,
    max_iterations: int | None = None,
    initial: tg.Tensor | None = None,
    preconditioner: Preconditioner | None = None,
) -> tg.Tensor:
    """Capture bounded matrix-free (preconditioned) conjugate gradient.

    The operator contract is symmetric positive-definite, as in scipy; the
    sugar layer does not attempt a symmetry proof. Four values (solution,
    residual, search direction, residual inner product) are carried through
    one ``gf_control.repeat`` or ``gf_control.while`` region, so the operator
    and optional preconditioner remain ordinary compiler-visible
    Tensor/MessagePassing programs.

    Pass ``iterations=k`` for a fixed ``gf_control.repeat`` or pass both
    ``tolerance=eps`` and ``max_iterations=k`` for a device-side bounded
    ``gf_control.while``. The convergence condition is the absolute Euclidean
    residual norm and never synchronizes a scalar through Python.
    """
    matvec, extent = _bind_operator(
        operator, graph=graph, field=field, src=src, dst=dst,
        edge=edge, params=params)
    if not isinstance(rhs, tg.Tensor) or rhs.ndim != 1:
        raise TypeError("cg rhs must be a rank-one tiga.Tensor")
    if rhs.dtype.kind != "float":
        raise TypeError("cg currently requires a real floating Tensor")
    if extent is not None and rhs.shape != (extent,):
        raise ValueError(f"cg rhs shape must be {(extent,)}, got {rhs.shape}")
    fixed = iterations is not None
    convergent = tolerance is not None or max_iterations is not None
    if fixed == convergent:
        raise ValueError(
            "cg requires exactly one stopping contract: iterations, or "
            "tolerance with max_iterations")
    if fixed:
        if not isinstance(iterations, int) or isinstance(iterations, bool):
            raise TypeError("cg iterations must be an integer")
        if iterations < 0:
            raise ValueError("cg iterations must be non-negative")
    else:
        if tolerance is None or max_iterations is None:
            raise ValueError(
                "tolerance-driven cg requires tolerance and max_iterations")
        if (not isinstance(max_iterations, int) or
                isinstance(max_iterations, bool)):
            raise TypeError("cg max_iterations must be an integer")
        if max_iterations < 0:
            raise ValueError("cg max_iterations must be non-negative")
    if preconditioner is not None and not callable(preconditioner):
        raise TypeError("cg preconditioner must be a Tensor -> Tensor callable")
    solution = tg.zeros_like(rhs) if initial is None else initial
    if (not isinstance(solution, tg.Tensor) or solution.shape != rhs.shape or
            solution.dtype is not rhs.dtype or solution.device != rhs.device):
        raise ValueError("cg initial value must match rhs shape, dtype, and device")

    def apply_preconditioner(value: tg.Tensor) -> tg.Tensor:
        if preconditioner is None:
            return value
        output = preconditioner(value)
        if (not isinstance(output, tg.Tensor) or output.shape != value.shape or
                output.dtype is not value.dtype or output.device != value.device):
            raise ValueError(
                "cg preconditioner must preserve residual shape, dtype, and device")
        return output

    residual = rhs - _checked_apply(matvec, solution, extent)
    preconditioned = apply_preconditioner(residual)
    residual_product = dot(residual, preconditioned)

    def checked(value: tg.Tensor) -> tg.Tensor:
        return _checked_apply(matvec, value, extent)

    if fixed:
        assert iterations is not None
        return _cg_fixed_loop(
            checked, apply_preconditioner, solution, residual, preconditioned,
            residual_product, iterations)
    assert tolerance is not None and max_iterations is not None
    if isinstance(tolerance, tg.Tensor):
        threshold = tolerance
        if (threshold.shape != () or threshold.dtype is not rhs.dtype or
                threshold.device != rhs.device):
            raise ValueError(
                "cg Tensor tolerance must be scalar and match rhs dtype/device")
    else:
        if not isinstance(tolerance, (int, float)) or isinstance(tolerance, bool):
            raise TypeError("cg tolerance must be a scalar Tensor or real number")
        if tolerance < 0:
            raise ValueError("cg tolerance must be non-negative")
        threshold = tg.tensor(tolerance, dtype=rhs.dtype, device=rhs.device)
    threshold_squared = threshold * threshold
    return _cg_tolerance_loop(
        checked, apply_preconditioner, solution, residual, preconditioned,
        residual_product, threshold_squared, max_iterations)


def bicgstab(
    operator,
    rhs: tg.Tensor,
    *,
    graph: tg.Graph | None = None,
    field: str | None = None,
    src: Mapping[str, tg.Tensor] | None = None,
    dst: Mapping[str, tg.Tensor] | None = None,
    edge: Mapping[str, tg.Tensor] | None = None,
    params: Mapping[str, object] | None = None,
    iterations: int | None = None,
    tolerance: tg.Tensor | float | None = None,
    max_iterations: int | None = None,
    initial: tg.Tensor | None = None,
) -> tg.Tensor:
    """Capture bounded matrix-free BiCGStab for nonsymmetric operators.

    Same stopping contracts as ``cg``: ``iterations=k`` for a fixed
    ``gf_control.repeat``, or ``tolerance=eps`` plus ``max_iterations=k`` for
    a bounded device-side ``gf_control.while``. Seven values (solution,
    residual, search direction, operator images, and the scalar recurrences)
    are carried through one control region.

    The convergence test runs at the top of each iteration, as in ``cg``.
    Landing exactly on the solution at a half step (residual surrogate
    ``s = 0``) divides by ``dot(A·s, A·s)`` — the same exact-breakdown hazard
    fixed-count CG has; a nonzero tolerance exits before it in practice.
    """
    matvec, extent = _bind_operator(
        operator, graph=graph, field=field, src=src, dst=dst,
        edge=edge, params=params)
    if not isinstance(rhs, tg.Tensor) or rhs.ndim != 1:
        raise TypeError("bicgstab rhs must be a rank-one tiga.Tensor")
    if rhs.dtype.kind != "float":
        raise TypeError("bicgstab currently requires a real floating Tensor")
    if extent is not None and rhs.shape != (extent,):
        raise ValueError(f"bicgstab rhs shape must be {(extent,)}, got {rhs.shape}")
    fixed = iterations is not None
    convergent = tolerance is not None or max_iterations is not None
    if fixed == convergent:
        raise ValueError(
            "bicgstab requires exactly one stopping contract: iterations, or "
            "tolerance with max_iterations")
    if fixed:
        if not isinstance(iterations, int) or isinstance(iterations, bool):
            raise TypeError("bicgstab iterations must be an integer")
        if iterations < 0:
            raise ValueError("bicgstab iterations must be non-negative")
    else:
        if tolerance is None or max_iterations is None:
            raise ValueError(
                "tolerance-driven bicgstab requires tolerance and max_iterations")
        if (not isinstance(max_iterations, int) or
                isinstance(max_iterations, bool)):
            raise TypeError("bicgstab max_iterations must be an integer")
        if max_iterations < 0:
            raise ValueError("bicgstab max_iterations must be non-negative")
    solution = tg.zeros_like(rhs) if initial is None else initial
    if (not isinstance(solution, tg.Tensor) or solution.shape != rhs.shape or
            solution.dtype is not rhs.dtype or solution.device != rhs.device):
        raise ValueError(
            "bicgstab initial value must match rhs shape, dtype, and device")

    residual = rhs - _checked_apply(matvec, solution, extent)
    shadow = residual                       # r̂ = r₀ stays constant
    direction = tg.zeros_like(rhs)
    velocity = tg.zeros_like(rhs)
    one = tg.tensor(1.0, dtype=rhs.dtype, device=rhs.device)
    rho = alpha = omega = one

    def checked(value: tg.Tensor) -> tg.Tensor:
        return _checked_apply(matvec, value, extent)

    if fixed:
        assert iterations is not None
        return _bicgstab_fixed_loop(
            checked, solution, residual, shadow, direction, velocity,
            rho, alpha, omega, iterations)
    assert tolerance is not None and max_iterations is not None
    if isinstance(tolerance, tg.Tensor):
        threshold = tolerance
        if (threshold.shape != () or threshold.dtype is not rhs.dtype or
                threshold.device != rhs.device):
            raise ValueError(
                "bicgstab Tensor tolerance must be scalar and match rhs")
    else:
        if not isinstance(tolerance, (int, float)) or isinstance(tolerance, bool):
            raise TypeError(
                "bicgstab tolerance must be a scalar Tensor or real number")
        if tolerance < 0:
            raise ValueError("bicgstab tolerance must be non-negative")
        threshold = tg.tensor(tolerance, dtype=rhs.dtype, device=rhs.device)
    threshold_squared = threshold * threshold
    return _bicgstab_tolerance_loop(
        checked, solution, residual, shadow, direction, velocity,
        rho, alpha, omega, threshold_squared, max_iterations)


def linear_solve(
    operator,
    rhs: tg.Tensor,
    *,
    method: str = "cg",
    graph: tg.Graph | None = None,
    field: str | None = None,
    src: Mapping[str, tg.Tensor] | None = None,
    dst: Mapping[str, tg.Tensor] | None = None,
    edge: Mapping[str, tg.Tensor] | None = None,
    params: Mapping[str, object] | None = None,
    iterations: int | None = None,
    tolerance: tg.Tensor | float | None = None,
    max_iterations: int | None = None,
    initial: tg.Tensor | None = None,
    preconditioner: Preconditioner | None = None,
    relaxation: tg.Tensor | float = 1.0,
) -> tg.Tensor:
    """Solve ``A·x = b`` matrix-free; ``method`` selects the iteration.

    - ``"cg"`` (default) — conjugate gradient; operator contract is symmetric
      positive-definite, optional ``preconditioner``.
    - ``"bicgstab"`` — BiCGStab for nonsymmetric operators.
    - ``"richardson"`` — fixed-count damped iteration; requires
      ``iterations=k`` and accepts ``relaxation``.

    All methods share the stopping contracts: ``iterations=k`` for a fixed
    ``gf_control.repeat``, or ``tolerance=eps`` with ``max_iterations=k`` for
    a bounded device-side ``gf_control.while``.
    """
    if method == "cg":
        return cg(
            operator, rhs, graph=graph, field=field, src=src, dst=dst,
            edge=edge, params=params, iterations=iterations,
            tolerance=tolerance, max_iterations=max_iterations,
            initial=initial, preconditioner=preconditioner)
    if method == "bicgstab":
        if preconditioner is not None:
            raise ValueError(
                "linear_solve method='bicgstab' takes no preconditioner in "
                "this sugar layer")
        return bicgstab(
            operator, rhs, graph=graph, field=field, src=src, dst=dst,
            edge=edge, params=params, iterations=iterations,
            tolerance=tolerance, max_iterations=max_iterations,
            initial=initial)
    if method == "richardson":
        if tolerance is not None or max_iterations is not None:
            raise ValueError(
                "linear_solve method='richardson' is fixed-count: pass "
                "iterations, not tolerance")
        if preconditioner is not None:
            raise ValueError(
                "linear_solve method='richardson' takes no preconditioner")
        return richardson(
            operator, rhs, graph=graph, field=field, src=src, dst=dst,
            edge=edge, params=params, iterations=iterations,
            relaxation=relaxation, initial=initial)
    raise ValueError(
        f"unknown linear_solve method {method!r}: "
        "choose 'cg', 'bicgstab' or 'richardson'")


__all__ = [
    "Matvec", "Preconditioner", "bicgstab", "cg", "dot", "linear_solve",
    "richardson", "vector_norm",
]

Picard 迭代的非线性扩散

是什么。 同样的 P1 网格、同样的停止契约——但算子现在是非线性的： 扩散系数依赖于解本身，k(u) = 1 + u²。方程 −∇·(k(u)∇u) = 1 描述热导率 随温度升高的稳态热传导。由于组装形式 F(u) = A(u)·u 随迭代变化， Krylov 方法不再直接适用；最简单的不动点迭代格式是 阻尼 Picard（非线性 Richardson）步，只需要算子像 F(u)——仍然是 一个 MessagePassing kernel，永远不是组装好的矩阵。阻尼 ω 必须让 F 的 Jacobi 矩阵满足 ω·λ_max < 2； 这里 λ_max ≈ 4·max(k)/h。

\[ F(u)_i \;=\; \frac{1}{h}\sum_{j} k_e\,(u_i - u_j) \;+\; \frac{\delta_i}{h}\Bigl(1 + \tfrac{u_i^2}{2}\Bigr)\,u_i, \qquad k_e \;=\; 1 + \frac{u_i^2 + u_j^2}{2} \]
\[ u \;\leftarrow\; u + \omega\,\bigl(b - F(u)\bigr) \]

每条边的 conductance k_e 对当前迭代两个端点的 k 取平均，因此 edge UDF 同时读取 src.u 和 dst.u；两个边界单元（端点处 δ_i = 1，内部为 0） 通过 node UDF 贡献一个对角项，消掉了纯差分 stencil 否则会有的常数零空间。 一个 nonlinear_solve(operator, b, omega=…) 驱动两种契约： iterations=k 捕获固定次数、可微分的 gf_control.repeat， tolerance=eps 捕获有界的设备端 gf_control.while——无论哪种都是 单个平坦循环，因为控制区域不支持嵌套。load_gradient() 沿展开的 repeat 对 sum(u) 关于载荷求导；迭代次数保持在个位数，因为展开 VJP 的编译时间随迭代数增长很快，测试用同一个映射的有限差分对梯度做了 交叉验证。

class NonlinearDiffusionApply(tg.MessagePassing):
    """Assemble F(u) = A(u)·u for -∇·(k(u)∇u) with k(u) = 1 + u².

    Edge (j -> i) sends the conductance-scaled difference
    k_e·(u_i - u_j)/h with the endpoint average k_e = 1 + (u_i² + u_j²)/2;
    the sum reducer assembles one nonlinear stiffness row per node. The two
    boundary elements (u fixed to 0) contribute the diagonal term
    (1 + u_i²/2)·u_i/h, added by the node UDF where dst.boundary = 1:

        F(u)_i = (1/h)·Σ_j k_e·(u_i - u_j) + (1/h)·δ_i·(1 + u_i²/2)·u_i,
        u₀ = u_{N+1} = 0.
    """

    reducer = tg.sum()

    def edge(self, src, dst, edge, inverse_spacing):
        del edge
        # src.u, dst.u: (N,) endpoint values of the current iterate;
        # the message is one scalar per edge.
        conductance = 1.0 + 0.5 * (src.u * src.u + dst.u * dst.u)  # k_e at u
        return conductance * (dst.u - src.u) * inverse_spacing

    def node(self, dst, incoming, inverse_spacing):
        # Boundary-element diagonal at the two end nodes (dst.boundary = 1);
        # k_b = 1 + u_i²/2 averages k(u_i) with k(0) = 1 at the fixed end.
        conductance = 1.0 + 0.5 * dst.u * dst.u
        return incoming + dst.boundary * conductance * dst.u * inverse_spacing


@tg.jit
def _picard_fixed_loop(apply, omega, state, rhs, iterations):
    for _ in range(iterations):
        state = state + omega * (rhs - apply(state))
    return state


@tg.jit
def _picard_tolerance_loop(
    apply, omega, state, rhs, threshold_squared, max_iterations,
):
    residual = rhs - apply(state)
    for _ in range(max_iterations):
        if dot(residual, residual) <= threshold_squared:
            break
        state = state + omega * residual
        residual = rhs - apply(state)
    return state


def nonlinear_solve(
    operator,
    rhs: tg.Tensor,
    *,
    omega: float,
    iterations: int | None = None,
    tolerance: float | None = 1.0e-6,
    max_iterations: int = 2_000,
    initial: tg.Tensor | None = None,
) -> tg.Tensor:
    """Solve F(u) = rhs by damped Picard (nonlinear Richardson) iteration.

    The fixed-point step u ← u + ω·(rhs − F(u)) runs in ONE flat loop:
    ``iterations=k`` captures a fixed ``gf_control.repeat`` (differentiable
    through ``tg.autograd.grad``; the VJP unrolls the loop, so keep k small
    when differentiating); ``tolerance=eps`` with ``max_iterations``
    captures a bounded device-side ``gf_control.while`` instead (forward-only
    — the while form has no VJP yet). The contracts are exclusive.
    """
    if not callable(operator):
        raise TypeError("operator must be a Tensor -> Tensor callable")
    if not isinstance(rhs, tg.Tensor) or rhs.ndim != 1:
        raise TypeError(
            "nonlinear_solve rhs must be a rank-one tiga.Tensor")
    if rhs.dtype.kind != "float":
        raise TypeError("nonlinear_solve requires a real floating Tensor")
    fixed = iterations is not None
    if fixed == (tolerance is not None):
        raise ValueError(
            "nonlinear_solve requires exactly one stopping contract: "
            "iterations, or tolerance with max_iterations")
    if fixed:
        if not isinstance(iterations, int) or isinstance(iterations, bool):
            raise TypeError("nonlinear_solve iterations must be an integer")
        if iterations < 0:
            raise ValueError("nonlinear_solve iterations must be non-negative")
    else:
        if not isinstance(tolerance, (int, float)) or isinstance(tolerance, bool):
            raise TypeError("nonlinear_solve tolerance must be a real number")
        if tolerance < 0:
            raise ValueError("nonlinear_solve tolerance must be non-negative")
        if (not isinstance(max_iterations, int) or
                isinstance(max_iterations, bool)):
            raise TypeError("nonlinear_solve max_iterations must be an integer")
        if max_iterations < 0:
            raise ValueError(
                "nonlinear_solve max_iterations must be non-negative")
    state = tg.zeros_like(rhs) if initial is None else initial
    if not isinstance(state, tg.Tensor) or state.shape != rhs.shape:
        raise ValueError("nonlinear_solve initial value must match rhs shape")
    if state.dtype is not rhs.dtype or state.device != rhs.device:
        raise ValueError(
            "nonlinear_solve initial value must match rhs dtype and device")
    if fixed:
        return _picard_fixed_loop(operator, omega, state, rhs, iterations)
    threshold = tg.tensor(
        tolerance * tolerance, dtype=rhs.dtype, device=rhs.device)
    return _picard_tolerance_loop(
        operator, omega, state, rhs, threshold, max_iterations)


def nonlinear_operator(interior_nodes: int):
    """P1 stencil grid + kernel for -∇·((1+u²)∇u) on the unit interval."""
    if interior_nodes < 1:
        raise ValueError("interior_nodes must be positive")
    spacing = 1.0 / (interior_nodes + 1)
    graph = tg.Graph.stencil((interior_nodes,), ((-1,), (0,), (1,)))
    flags = (                                   # one entry per boundary element
        [2.0] if interior_nodes == 1
        else [1.0] + [0.0] * (interior_nodes - 2) + [1.0])
    boundary = tg.tensor(flags, dtype=tg.float32)  # (N,)
    return NonlinearDiffusionApply(), graph, spacing, boundary


def solve(
    interior_nodes: int = 8,
    *,
    omega: float = 0.05,
    iterations: int | None = None,
    tolerance: float | None = 1.0e-6,
    max_iterations: int = 2_000,
):
    """Solve F(u) = h·1 for the nonlinear diffusion operator.

    Default is a residual-driven ``gf_control.while``; pass ``iterations=k``
    for a fixed ``gf_control.repeat`` loop instead. ω must stay below
    2/λ_max of the Jacobian — λ_max ≈ 4·max(k)/h here, so ω = 0.05 is
    stable for both grid sizes used below.
    """
    kernel, graph, spacing, boundary = nonlinear_operator(interior_nodes)
    load = tg.tensor(
        [spacing] * interior_nodes, dtype=tg.float32, requires_grad=True)  # (N,)

    def apply(state):  # F(u): (N,) -> (N,), conductance from current iterate
        return kernel(
            graph=graph,
            src={"u": state},
            dst={"u": state, "boundary": boundary},
            inverse_spacing=1.0 / spacing,
        )

    if iterations is not None:  # fixed repeat and bounded while are exclusive
        tolerance = None
    solution = nonlinear_solve(
        apply,
        load,
        omega=omega,
        iterations=iterations,
        tolerance=tolerance,
        max_iterations=max_iterations,
    )  # (N,)
    residual = load - apply(solution)                        # (N,)
    error = max(abs(value) for value in residual.tolist())
    return solution, residual, error
完整源码：examples/nonlinear_solve.py（可直接运行） 
"""Matrix-free 1D nonlinear diffusion solve through MessagePassing and Picard."""

from __future__ import annotations

import tiga as tg
from solvers import dot


class NonlinearDiffusionApply(tg.MessagePassing):
    """Assemble F(u) = A(u)·u for -∇·(k(u)∇u) with k(u) = 1 + u².

    Edge (j -> i) sends the conductance-scaled difference
    k_e·(u_i - u_j)/h with the endpoint average k_e = 1 + (u_i² + u_j²)/2;
    the sum reducer assembles one nonlinear stiffness row per node. The two
    boundary elements (u fixed to 0) contribute the diagonal term
    (1 + u_i²/2)·u_i/h, added by the node UDF where dst.boundary = 1:

        F(u)_i = (1/h)·Σ_j k_e·(u_i - u_j) + (1/h)·δ_i·(1 + u_i²/2)·u_i,
        u₀ = u_{N+1} = 0.
    """

    reducer = tg.sum()

    def edge(self, src, dst, edge, inverse_spacing):
        del edge
        # src.u, dst.u: (N,) endpoint values of the current iterate;
        # the message is one scalar per edge.
        conductance = 1.0 + 0.5 * (src.u * src.u + dst.u * dst.u)  # k_e at u
        return conductance * (dst.u - src.u) * inverse_spacing

    def node(self, dst, incoming, inverse_spacing):
        # Boundary-element diagonal at the two end nodes (dst.boundary = 1);
        # k_b = 1 + u_i²/2 averages k(u_i) with k(0) = 1 at the fixed end.
        conductance = 1.0 + 0.5 * dst.u * dst.u
        return incoming + dst.boundary * conductance * dst.u * inverse_spacing


@tg.jit
def _picard_fixed_loop(apply, omega, state, rhs, iterations):
    for _ in range(iterations):
        state = state + omega * (rhs - apply(state))
    return state


@tg.jit
def _picard_tolerance_loop(
    apply, omega, state, rhs, threshold_squared, max_iterations,
):
    residual = rhs - apply(state)
    for _ in range(max_iterations):
        if dot(residual, residual) <= threshold_squared:
            break
        state = state + omega * residual
        residual = rhs - apply(state)
    return state


def nonlinear_solve(
    operator,
    rhs: tg.Tensor,
    *,
    omega: float,
    iterations: int | None = None,
    tolerance: float | None = 1.0e-6,
    max_iterations: int = 2_000,
    initial: tg.Tensor | None = None,
) -> tg.Tensor:
    """Solve F(u) = rhs by damped Picard (nonlinear Richardson) iteration.

    The fixed-point step u ← u + ω·(rhs − F(u)) runs in ONE flat loop:
    ``iterations=k`` captures a fixed ``gf_control.repeat`` (differentiable
    through ``tg.autograd.grad``; the VJP unrolls the loop, so keep k small
    when differentiating); ``tolerance=eps`` with ``max_iterations``
    captures a bounded device-side ``gf_control.while`` instead (forward-only
    — the while form has no VJP yet). The contracts are exclusive.
    """
    if not callable(operator):
        raise TypeError("operator must be a Tensor -> Tensor callable")
    if not isinstance(rhs, tg.Tensor) or rhs.ndim != 1:
        raise TypeError(
            "nonlinear_solve rhs must be a rank-one tiga.Tensor")
    if rhs.dtype.kind != "float":
        raise TypeError("nonlinear_solve requires a real floating Tensor")
    fixed = iterations is not None
    if fixed == (tolerance is not None):
        raise ValueError(
            "nonlinear_solve requires exactly one stopping contract: "
            "iterations, or tolerance with max_iterations")
    if fixed:
        if not isinstance(iterations, int) or isinstance(iterations, bool):
            raise TypeError("nonlinear_solve iterations must be an integer")
        if iterations < 0:
            raise ValueError("nonlinear_solve iterations must be non-negative")
    else:
        if not isinstance(tolerance, (int, float)) or isinstance(tolerance, bool):
            raise TypeError("nonlinear_solve tolerance must be a real number")
        if tolerance < 0:
            raise ValueError("nonlinear_solve tolerance must be non-negative")
        if (not isinstance(max_iterations, int) or
                isinstance(max_iterations, bool)):
            raise TypeError("nonlinear_solve max_iterations must be an integer")
        if max_iterations < 0:
            raise ValueError(
                "nonlinear_solve max_iterations must be non-negative")
    state = tg.zeros_like(rhs) if initial is None else initial
    if not isinstance(state, tg.Tensor) or state.shape != rhs.shape:
        raise ValueError("nonlinear_solve initial value must match rhs shape")
    if state.dtype is not rhs.dtype or state.device != rhs.device:
        raise ValueError(
            "nonlinear_solve initial value must match rhs dtype and device")
    if fixed:
        return _picard_fixed_loop(operator, omega, state, rhs, iterations)
    threshold = tg.tensor(
        tolerance * tolerance, dtype=rhs.dtype, device=rhs.device)
    return _picard_tolerance_loop(
        operator, omega, state, rhs, threshold, max_iterations)


def nonlinear_operator(interior_nodes: int):
    """P1 stencil grid + kernel for -∇·((1+u²)∇u) on the unit interval."""
    if interior_nodes < 1:
        raise ValueError("interior_nodes must be positive")
    spacing = 1.0 / (interior_nodes + 1)
    graph = tg.Graph.stencil((interior_nodes,), ((-1,), (0,), (1,)))
    flags = (                                   # one entry per boundary element
        [2.0] if interior_nodes == 1
        else [1.0] + [0.0] * (interior_nodes - 2) + [1.0])
    boundary = tg.tensor(flags, dtype=tg.float32)  # (N,)
    return NonlinearDiffusionApply(), graph, spacing, boundary


def solve(
    interior_nodes: int = 8,
    *,
    omega: float = 0.05,
    iterations: int | None = None,
    tolerance: float | None = 1.0e-6,
    max_iterations: int = 2_000,
):
    """Solve F(u) = h·1 for the nonlinear diffusion operator.

    Default is a residual-driven ``gf_control.while``; pass ``iterations=k``
    for a fixed ``gf_control.repeat`` loop instead. ω must stay below
    2/λ_max of the Jacobian — λ_max ≈ 4·max(k)/h here, so ω = 0.05 is
    stable for both grid sizes used below.
    """
    kernel, graph, spacing, boundary = nonlinear_operator(interior_nodes)
    load = tg.tensor(
        [spacing] * interior_nodes, dtype=tg.float32, requires_grad=True)  # (N,)

    def apply(state):  # F(u): (N,) -> (N,), conductance from current iterate
        return kernel(
            graph=graph,
            src={"u": state},
            dst={"u": state, "boundary": boundary},
            inverse_spacing=1.0 / spacing,
        )

    if iterations is not None:  # fixed repeat and bounded while are exclusive
        tolerance = None
    solution = nonlinear_solve(
        apply,
        load,
        omega=omega,
        iterations=iterations,
        tolerance=tolerance,
        max_iterations=max_iterations,
    )  # (N,)
    residual = load - apply(solution)                        # (N,)
    error = max(abs(value) for value in residual.tolist())
    return solution, residual, error


def load_gradient(
    interior_nodes: int = 4, *, omega: float = 0.05, iterations: int = 3,
):
    """Differentiate sum(u) w.r.t. the load through the fixed Picard repeat.

    The fixed-repeat VJP unrolls the loop, so its compile time grows quickly
    with the iteration count — keep it in the low single digits (the
    forward-only drivers above run hundreds of iterations cheaply).
    """
    kernel, graph, spacing, boundary = nonlinear_operator(interior_nodes)
    load = tg.tensor(
        [spacing] * interior_nodes, dtype=tg.float32, requires_grad=True)  # (N,)

    def apply(state):  # F(u): (N,) -> (N,)
        return kernel(
            graph=graph,
            src={"u": state},
            dst={"u": state, "boundary": boundary},
            inverse_spacing=1.0 / spacing,
        )

    solution = nonlinear_solve(
        apply, load, omega=omega, iterations=iterations, tolerance=None)  # (N,)
    return tg.autograd.grad(solution.sum(), load)              # (N,)


# Inspect: solve()[0].mlir() (bounded gf_control.while loop),
#   solve(iterations=200)[0].mlir() (fixed gf_control.repeat loop)

if __name__ == "__main__":
    solve()
    load_gradient()
本机编译产物（Ryzen 7 255 · RTX 5070 Ti 实测） 
### NonlinearDiffusionApply
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=3

### NonlinearDiffusionApply
backend: cpu:0
provider: tiga-runtime
lowering: gf-tensor-relation-autograd
passes: bind-static-csr-snapshot, capture-message-passing-udf, analyze-reducer-algebra, lower-csr-to-gather-segment, fuse-edge-node-regions
remark: [planning] edge/node UDFs remain in the differentiable Tensor DAG
remark: [planning] gf-tensor-vjp generates CSR gather/segment-sum adjoints
remark: [planning] reducer lowering: builtin-additive-state
variant cache: hits=0, misses=1