tensorplay@1.0.0.dev20260909

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.

Creates the CustomOpDef if absent (kernels are then attached with impl() or Library("ns", "IMPL").impl). tags is Like Library.define, a full "ns::op(Tensor) -> Tensor" string may be pasted as qualname.

dispatch_key accepts "cpu"/"cuda"/"default" and the composite spellings. Raises LookupError when nothing usable is registered (a disabled concrete kernel counts as absent, matching CustomOpDef.set_kernel_enabled() visibility).

types accepts concrete devices ("CPU"/"CUDA") or composite spellings (CompositeExplicitAutograd → the device-agnostic slot).

Produces "ns::op(Tensor self, SymInt n, Tensor(a!) out) -> Tensor" alias-annotation (<type>(<letter>!)) marker. Parameters without annotations are treated as tensors; *args/**kwargs are skipped.

Runs each selected check and reports failures keyed by test name:

• test_schema: undeclared inputs are left unmutated and no output aliases an input storage (declared-mutation direction is trusted, matching TensorPlay’s declaration-driven fusion barriers).

• test_autograd_registration: gradients reach every floating input with matching shapes. TensorPlay composes Python kernels implicitly (CompositeImplicitAutograd semantics), so a missing explicit formula is legal — this check catches kernels that break the autograd graph or drop gradients.

• test_faketensor: the fake kernel reproduces the real outputs’ metadata.

• test_aot_dispatch_dynamic: capture + execution reproduce the

Returns the failure mapping; empty means all checks passed.

Accepts a CustomOpDef or a qualified operator name. device_types=None or an empty iterable means the device-agnostic slot; composite spellings (Composite…) map there too. Usable directly or as a decorator.

The registered kernel(s) must launch their Triton kernels through wrap_triton() and only mutate arguments listed in mutates_args. Under tensorplay.compile the whole operator is captured as a single opaque node — the compiler never traces into the Triton launches. device_types is a

Mark a Triton kernel as launchable from within a triton_op.

Accepts a triton.runtime.jit.JITFunction (the @triton.jit result) or any grid-indexable launcher; idempotent on wrappers.