Torch Loop

torch_loop

Explicit, bounded CPU Torch compilation for sequential solver loops.

Importing this module does not import Torch. Constructing an adapter opts into native Inductor compilation; ordinary solver calls retain their eager defaults.

torch_compile_loop_adapter(*, max_steps=128, max_elements=65536)

Construct an explicitly registered, bounded sequential CPU adapter.

Supports Tensor/None trees in tuples, lists and dictionaries, homogeneous CPU float32/64, positive static input lengths and first-order autograd. The sequential loop is statically captured with fullgraph Inductor; graph breaks, disabled compilation and unsupported inputs raise. Checkpointing and GPU execution are not part of this adapter's qualified contract.

Parameters:

Name	Type	Description	Default
max_steps	int	Largest positive scan length accepted by this adapter.	128
max_elements	int	Bound on input, carry and stacked output tensor elements.	65536

Returns:

Type	Description
LoopAdapter	An unregistered adapter; register it for array_api_compat.torch.

Raises:

Type	Description
ValueError	If bounds or Torch version are unsupported.
ImportError	If the optional Torch dependency is absent.

Source code in src/op_engine/torch_loop.py

def torch_compile_loop_adapter(
    *, max_steps: int = 128, max_elements: int = 65_536
) -> LoopAdapter:
    """Construct an explicitly registered, bounded sequential CPU adapter.

    Supports Tensor/None trees in tuples, lists and dictionaries, homogeneous
    CPU float32/64, positive static input lengths and first-order autograd.
    The sequential loop is statically captured with fullgraph Inductor; graph
    breaks, disabled compilation and unsupported inputs raise. Checkpointing
    and GPU execution are not part of this adapter's qualified contract.

    Args:
        max_steps: Largest positive scan length accepted by this adapter.
        max_elements: Bound on input, carry and stacked output tensor elements.

    Returns:
        An unregistered adapter; register it for ``array_api_compat.torch``.

    Raises:
        ValueError: If bounds or Torch version are unsupported.
        ImportError: If the optional Torch dependency is absent.
    """  # noqa: DOC502 - optional import propagates ImportError
    if type(max_steps) is not int or not 1 <= max_steps <= 128:
        msg = "max_steps must be an integer in [1, 128]"
        raise ValueError(msg)
    if type(max_elements) is not int or not 1 <= max_elements <= 1_048_576:
        msg = "max_elements must be an integer in [1, 1048576]"
        raise ValueError(msg)
    torch = import_module("torch")
    version = tuple(int(part) for part in torch.__version__.split(".")[:2])
    if version < (2, 14):
        msg = "Compiled loops require Torch >=2.14; eager Torch is unchanged"
        raise ValueError(msg)
    driver = _TorchScan(
        torch, import_module("torch.utils._pytree"), max_steps, max_elements
    )
    # Torch specialization caches use code identity. Isolate each explicit
    # adapter's recompilation budget without retaining model inputs globally.
    unique = FunctionType(_TorchScan.run.__code__.replace(), globals())
    compiled = torch.compile(
        MethodType(unique, driver), backend="inductor", fullgraph=True, dynamic=False
    )

    def scan(body: LoopBody, carry: Any, inputs: Any) -> tuple[Any, Any]:  # noqa: ANN401
        if torch.compiler.is_compiling():
            return driver.run(body, carry, inputs)
        driver.check_inputs(carry, inputs)
        return cast("tuple[Any, Any]", compiled(body, carry, inputs))

    return LoopAdapter(scan=cast("ScanFunction", scan))