tvm_ffi.convert_func

Contents

tvm_ffi.convert_func#

tvm_ffi.convert_func(pyfunc, tensor_cls=None)[source]#

Convert a Python callable to an FFI Function.

This is the callable-specific sibling of tvm_ffi.convert(). It accepts one extra argument, tensor_cls, that lets the caller specify how tensor arguments should be delivered to the Python callable when the resulting Function is invoked from C++. tvm_ffi.convert() has no such knob — it always produces a Function whose callback receives tvm_ffi.Tensor for tensor args.

Parameters:

  • pyfunc (Callable) – The Python callable to wrap.

  • tensor_cls (type, optional) –

    The class whose instances the callback should receive for tensor args. The class must expose a __dlpack_c_exchange_api__ PyCapsule; its capsule is threaded into the callback closure so tensor args are converted at the C level (via the DLPack exchange API) before the Python callback body runs — this is significantly faster than calling torch.from_dlpack(x) (or equivalent) inside the callback. Raises TypeError if tensor_cls does not expose the attribute.

    When tensor_cls is None, convert_func behaves like the callable branch of tvm_ffi.convert().

Return type:

Any

Returns:

Function – The wrapped FFI function.

Examples

import torch
import tvm_ffi

# Without tensor_cls: same as tvm_ffi.convert(pyfunc) — the callback
# receives tvm_ffi.Tensor for tensor args.
f = tvm_ffi.convert_func(lambda x: x + 1)
assert isinstance(f, tvm_ffi.Function)

# With tensor_cls=torch.Tensor: the callback receives torch.Tensor
# directly; the DLPack conversion happens in C before the body runs.
def callback(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    return a + b


g = tvm_ffi.convert_func(callback, tensor_cls=torch.Tensor)

See also

tvm_ffi.convert()

Generic value-to-FFI conversion. Use this when you don’t need to specify tensor_cls.

Contents