Why a defined custom-op path matters
The hard part of a custom operator is not putting a function in a namespace. It is preserving one semantic contract across eager execution, autograd, graph capture, and backend execution. TensorPlay splits that path into registration, dispatch, differentiation, and native lowering boundaries that can be inspected independently.
tensorplay.library exposes custom_op, triton_op, wrap_triton, register_kernel, register_fake, and register_autograd. An operator can start with a Python implementation for numerical checks and then gain a device-specific kernel without changing the model above it.
Three layers: registration, execution, and autograd
- Registration: give the operator a namespace and schema, then provide a default implementation and device-specific kernels.
- Execution: eager mode selects the available kernel; unsupported capability must follow an explicit fallback contract.
- Autograd: register fake metadata and a backward rule so capture and differentiation do not depend on a hidden Python side effect.
@library.custom_op("demo::scaled", mutates_args=())
def scaled(x):
return tp.mul(x, 2.0)
@scaled.register_autograd
def backward(ctx, grad_out):
return (tp.mul(grad_out, 2.0),)Triton: eager execution with an explicit capture contract
Triton is useful when a kernel needs a device-level implementation, but a callable that works in eager mode is not automatically graph-safe. The repository tests treat the Triton path as an opaque node during capture and make the boundary visible to the compiler instead of pretending the Python body can always be traced.
That distinction matters for debugging: first compare the eager kernel with the default implementation, then capture the same call, inspect the graph node and codegen marker, and finally run the compiled path on the target device. A missing CUDA environment should be an explicit skipped capability, not an unreported claim.
TVM: lower pointwise chains to TIR
The TVM backend targets a narrower and more useful contract: pointwise TensorPlay graphs can be lowered to TIR, compiled, and compared with the ordinary path. The backend tests cover parity, shape changes, training regions, and fallback, which gives the lowering step a concrete correctness surface.
This is not a claim that every graph lowers to TVM. It is a documented route for the subset the backend understands, with a fallback that remains observable. Keeping the supported subset small makes it possible to add operators without losing a reliable reference path.
Verification: check values, graph, and executor
- Compare eager and lowered outputs with fixed shapes, dtypes, and tolerances.
- Check that the graph contains the expected custom-op node and does not leak an accidental internal implementation.
- Check the codegen marker to confirm whether the result came from native, Triton, or TVM execution.
- Exercise unsupported backends and shape changes to verify fallback behavior instead of treating fallback as a performance result.