tensorplay.library.custom_op

tensorplay.library.custom_op(name: str, fn: Callable[[...], Any] | None = None, /, *, mutates_args: Sequence[str] = (), device_types: Any = None, schema: str | None = None) → Callable[[Callable[[...], Any]], CustomOpDef][source]

Example:

@tensorplay.library.custom_op("mylib::weighted_sum", mutates_args=())
def weighted_sum(x, weight):
    return (x * weight).sum()

# Optional extra kernels per device:
@weighted_sum.register_kernel("cuda")
def _(x, weight): ...

fn may also be passed positionally (custom_op("mylib::op", my_fn, mutates_args=())), matching operator’s default kernel for the advertised device_types (every device when omitted).

Parameters:

• name – Qualified "namespace::name" identifier.

• fn – Operator body; omit to use the return value as a decorator.

• mutates_args – Names of arguments the kernel mutates in place. Compile-time fusion treats these as barriers regardless of the value; eager execution trusts the declaration.

• device_types – Restriction advertised to users at definition time. Kernels are selected per call from whatever was registered via CustomOpDef.register_kernel().

• schema – Optional schema string kept for introspection and opcheck() (TensorPlay models no schema grammar).

Returns:

A decorator producing a callable CustomOpDef.