tvm_ffi.structural_walk

tvm_ffi.structural_walk(root, callbacks=(), with_def_region_kind=(), order='pre')

Walk a value structurally and invoke the first matching typed callback.

Parameters:

• root (Any) – Root value to traverse.

• callbacks (tuple | Sequence | Callable, default: ()) –

  Normal callbacks. These callbacks receive one argument: value. Callback entries are tried in order.

  May be one of:

  • A single callback, used as a typing.Any catch-all.

  • A (type, callback) entry.

  • A grouped ((type1, type2, ...), callback) entry.

  • A sequence of entries.

  Types may be builtins, registered FFI object classes, or typing.Any/object as a catch-all.

• with_def_region_kind (tuple | Sequence | Callable, default: ()) – Def-region-aware callbacks. These callbacks receive two arguments: (value, def_region_kind). They accept the same callback entry forms as callbacks.

• order (str | WalkOrder, default: "pre") – "pre"/WalkOrder.PREORDER to invoke callbacks before children, or "post"/WalkOrder.POSTORDER to invoke callbacks after children.

Return type:

VisitInterrupt | None

Returns:

result – None if traversal completed, otherwise a VisitInterrupt returned by a callback.

Examples

visited = []

uses = []
result = tvm_ffi.structural_walk(
    node,
    ((int, float), lambda value: visited.append(("leaf", value))),
    with_def_region_kind=(
        Var,
        lambda var, kind: (
            uses.append(var) if kind == tvm_ffi.DefRegionKind.NONE else None
        ),
    ),
)