Skip to content

Linear solvers and control flow

Define an operator that maps one vector to another, then pass it to a solver. Start with the three-step Poisson example below; dynamic neighborhoods, iteration algorithms and nonlinear equations follow afterward.

Native solver tensors and Torch interoperability

A kernel call accepts either — src/dst/edge fields may be torch tensors directly, without a wrapper. The advanced solver drivers on this page are different: linear_solve / nonlinear_solve capture the iteration loop as device-side control flow (gf_control.repeat / gf_control.while), so their vectors must be tg.Tensor. Wrap torch storage at the boundary with tg.from_torch(x) when entering native expression capture. Only a value that still owns Torch storage supports zero-copy .to_torch(); a native solver result requires .to_torch(copy=True) and contiguous storage. Sharing storage does not bridge native gradients into PyTorch autograd.

Solve Poisson with neighbor sums

Solve the Poisson equation on a line, with both endpoints fixed to zero and seven unknown interior nodes.

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

Define only how to compute A(u): each node takes twice its own value minus its two neighbors, then divides by the grid spacing. The solver owns the iteration loop; no handwritten CG loop is needed.

The complete computation follows. linear_solve is the repository helper examples/solvers.py, not a built-in 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))

Run 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

“Matrix-free” means no assembled stiffness matrix: Graph.stencil supplies adjacency, edge() negates the neighbor value, and node() adds twice the center. The graph still has CSR topology; there is no separate stiffness matrix or per-edge stiffness array.

This example uses tg.Tensor because the current solver helper captures native control flow, not because ordinary message passing requires replacing Torch tensors.

Why FEM, and how are boundaries handled?

Piecewise-linear P1 finite elements on a uniform 1-D grid produce the following equation. h is the spacing and b is the load vector for a unit source:

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

The graph contains only interior nodes and omits out-of-grid neighbors. Those boundary values are zero, but the diagonal coefficient in node() remains 2. Do not replace it with the valid-neighbor count.

One-dimensional P1 stencil: interior nodes read their two neighbors; boundary values remain zero.

The exact solution checks the nodal values:

\[ u_{\mathrm{exact}}(x)=\frac{x(1-x)}{2} \]
Stopping, fixed iterations and gradients

tolerance=1e-6 checks the absolute residual norm; max_iterations=32 bounds the loop. Reaching the limit does not guarantee convergence, so the example also checks the analytic-solution error.

Fixed-count iteration and load gradients remain in the advanced fem_poisson.py example. See linear solvers for control-flow and differentiation contracts.

Dynamic-radius matrix-free linear solve

What it is. Here there is no mesh file at all — the graph is generated from data. Given N points in the plane, any two distinct points at distance at most r become neighbors (a radius graph). On that graph the example solves (m·I + L) u = b, where L is the graph Laplacian: for each point, the sum over its neighbors of the difference between the point's own value and the neighbor's value. The mass shift m = 1 keeps the operator positive definite, and the right-hand side is the periodic pattern b_i = 1 + (i mod 3).

Radius relation: distinct points at distance at most r become neighbors; the operator at point i is m·u_i plus the sum over neighbors of 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) \]

The example places the points on a line with spacing h and sets r = 1.01·h, so each point's neighborhood is exactly its two immediate neighbors and the generated relation is a path.

The solver and operator interfaces are unchanged from the FEM case: the topology is produced by Graph.radius, and relation realization and reuse are provider decisions below the MessagePassing UDF. CG captures its convergence test and four carried states in one gf_control.while; no adjacency or sparse matrix is written by the user.

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
Full source: examples/meshfree_linear_solve.py (runs as-is) 
"""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()
Measured compilation artifacts (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

Shared solver sugar: linear_solve

What it is. Both examples above solve A·u = b by stationary iteration: start from a guess and repeatedly apply A to shrink the residual r = b − A·u. Three methods are provided, selected by one method= argument:

method= Algorithm Operator contract
"cg" (default) conjugate gradient symmetric positive-definite
"bicgstab" BiCGStab nonsymmetric
"richardson" damped fixed-step iteration any, slowest

Richardson takes the naive correction step u ← u + ω·r. Conjugate gradient does better when A is symmetric positive-definite: it walks a sequence of A-orthogonal search directions, carrying four states between iterations — the solution u, the residual r, the search direction p, and the inner product ρ = ⟨r, z⟩, where z is the (optionally preconditioned) residual. BiCGStab drops the symmetry requirement at the price of two operator applications per step, carrying additionally the shadow residual r̂ = r₀, the operator images v = A·p and t = A·s, and the scalars ρ, α, ω.

Stationary solver loop: operator, right-hand side and stopping contract enter a loop body captured once; states u, r, p and ρ are carried between iterations

\[ 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 composes these from gf_control primitives, with the loops written as natural Python for/while under @tg.jit. Each example wraps its operator construction in a local solve(); the solver itself is one linear_solve(operator, rhs, method="cg", …) call — there is nothing else named solve. The module is deliberately not part of the core package — it is grammar sugar, not a primitive.

examples/solvers.py (full source) 
"""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",
]

Nonlinear diffusion by Picard iteration

What it is. Same P1 grid, same stopping contracts — but the operator is now nonlinear: the diffusion coefficient depends on the solution itself, k(u) = 1 + u². The equation −∇·(k(u)∇u) = 1 is the steady state of heat flow whose conductivity grows with temperature. Because the assembled form F(u) = A(u)·u changes with the iterate, Krylov methods no longer apply directly; the simplest fixed-point scheme is the damped Picard (nonlinear Richardson) step, which only ever needs the operator image F(u) — still one MessagePassing kernel, never an assembled matrix. The damping ω must keep ω·λ_max of the Jacobian of F below 2; here λ_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) \]

The per-edge conductance k_e averages k at the two endpoints of the current iterate, so the edge UDF reads both src.u and dst.u; the two boundary elements (δ_i = 1 at the end nodes, 0 inside) contribute a diagonal term through the node UDF, which removes the constant null mode the pure difference stencil would otherwise have. One nonlinear_solve(operator, b, omega=…) drives both contracts: iterations=k captures a fixed, differentiable gf_control.repeat, and tolerance=eps captures a bounded device-side gf_control.while — one flat loop either way, since control regions do not nest. load_gradient() differentiates sum(u) with respect to the load through the unrolled repeat; the count stays in the low single digits because the unrolled VJP's compile time grows quickly with iterations, and the test cross-checks the gradient against a finite difference of the same map.

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
Full source: examples/nonlinear_solve.py (runs as-is) 
"""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()
Measured compilation artifacts (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