# Source code for tensorplay.library

Source: https://www.tensorplay.cn/docs/_modules/tensorplay/library.html

```
"""Custom operator registration and compiler integration.

Four public layers are provided:

1. :func:`custom_op` / the :class:`Library` class register opaque custom
   operators with device-specific kernels, optional fake (meta) kernels,
   an autograd formula (:meth:`CustomOpDef.register_autograd`), a vector-
   map formula (:meth:`CustomOpDef.register_vmap`) and an autocast cast
   rule (:meth:`CustomOpDef.register_autocast`).
2. :func:`triton_op` + :func:`wrap_triton` and :func:`tile_lang_op` +
   :func:`wrap_tilelang` integrate user-written Triton / TileLang kernels:
   such operators behave like any other custom operator in eager mode and
   are captured as one opaque graph node by ``tensorplay.compile`` —
   preserving the fusion boundary through compilation.
3. Registered operators compose with every compiler backend: during capture
   a call whose arguments are symbolic records a single ``call_function``
   node targeting the :class:`CustomOpDef`; backends that cannot lower it
   treat the node as a barrier and fall back to the interpreter, which
   dispatches to the registered kernel.
   :func:`define`/:func:`impl`/:func:`impl_abstract`, :func:`infer_schema`,
   :func:`get_kernel`, :meth:`CustomOpDef.set_kernel_enabled` and the
   validation harness :func:`opcheck`.

``"namespace::name"`` string.  Schemas are not modeled by a C++ dispatcher;
the optional ``schema=`` strings and :func:`infer_schema` output attach to
the operator for introspection, documentation and ``opcheck``.

``OpOverload``/``overload``/``deprecated``/``fallthrough_kernel``/
``NAMELESS_SCHEMA`` and ``get_ctx`` are absent, ``Library.fallback`` raises
``NotImplementedError``, and re-registering a kernel replaces the previous
one instead of raising (hot swaps for interactive sessions).
"""

from __future__ import annotations

import contextlib
import inspect
import threading
import typing
import warnings
from collections.abc import Callable, Sequence
from typing import Any, Iterable

import tensorplay
from .graph import (
    GraphCaptureError,
    capture_call as _capture_call,
    capturing as _capturing,
    _iter_proxies as _walk_proxies,
)

__all__ = [
    "Library",
    "CustomOpDef",
    "custom_op",
    "triton_op",
    "tile_lang_op",
    "wrap_triton",
    "wrap_tilelang",
    "register_kernel",
    "register_fake",
    "register_autograd",
    "register_vmap",
    "register_autocast",
    "define",
    "impl",
    "impl_abstract",
    "infer_schema",
    "opcheck",
    "get_kernel",
    "get_op",
    "has_op",
]

_LOCK = threading.RLock()
# "ns::op" -> CustomOpDef; namespaces that already own a DEF Library.
_OP_REGISTRY: dict[str, "CustomOpDef"] = {}
_DEFINED_LIBRARY_NAMESPACES: set[str] = set()

import tensorplay._C._autograd as _autograd_c

# Hot-path aliases resolved once at import (this module is imported last by
# tensorplay/__init__, so every attribute below already exists).  The gate
# reads the grad mode through the native binding directly: the public
# wrapper only adds a Python frame around the same thread-local probe.
_is_grad_enabled = _autograd_c.is_grad_enabled
# One native scan over the call arguments replaces the full graph-node
# build when no input needs a gradient.  Older extensions without the
# binding keep the unconditional apply path.
_any_requires_grad = getattr(_autograd_c, "_any_requires_grad", None)
if _any_requires_grad is None:
    def _any_requires_grad(*args, **kwargs):
        return True
# One-crossing native apply for registered custom operators: grad gate,
# requires-grad scan, graph node, kernel call and context wiring all happen
# inside the extension.  Extensions without the binding fall back to
# Function.apply.
_CUSTOM_OP_AUTOGRAD_APPLY = getattr(
    _autograd_c, "custom_op_autograd_apply", None)
_CONTEXT_CLS: Any = None
# Native first-tensor device scan and the O(1) dispatch-mode stack probe;
# both fall back to the Python walks when the extension predates them.
_fast_device_key = getattr(tensorplay._C, "_first_device_key", None)
_len_dispatch_modes = getattr(tensorplay._C, "_len_dispatch_mode", None)

# Profiler session gate for automatic op-span emission.  Resolved lazily:
# the bridge function exists whenever the compiled extension ships it, and
# calling it is one atomic load.  Module-level None until the first lookup
# keeps attribute access off the hot path.
_profiler_is_active_fn = None
_profiler_probe_done = False

# Span source tag: distinguishes custom-op spans from hand-written
# record_function annotations in exported traces.
_SPAN_SRC_CUSTOM_OP = "custom_op"

def _profiler_active() -> bool:
    """True while a profiling session records on this process.

    One atomic load via the compiled bridge; the binding object is cached
    after the first call.  Falls back to False when the bridge is absent
    (older extension), keeping the wrapper callable.
    """
    global _profiler_is_active_fn, _profiler_probe_done
    if not _profiler_probe_done:
        _profiler_is_active_fn = getattr(
            tensorplay._C, "_profiler_is_active", None
        )
        _profiler_probe_done = True
    return _profiler_is_active_fn is not None and _profiler_is_active_fn()

# Live profiling sessions (nesting counted; the profiler start/stop hooks
# call _note_profiling_start/_note_profiling_stop).  The eager paths read
# this module global: one dict-free global load per custom-op call, so an
# inactive session costs nothing and an active one emits an op span.
_profiling_sessions = 0

def _note_profiling_start() -> None:
    global _profiling_sessions
    _profiling_sessions += 1

def _note_profiling_stop() -> None:
    global _profiling_sessions
    if _profiling_sessions > 0:
        _profiling_sessions -= 1

_COMPOSITE_KEYS = frozenset(
    {"CompositeExplicitAutograd", "CompositeImplicitAutograd"}
)
_COMPOSITE_KEYS_LOWERED = frozenset(key.lower() for key in _COMPOSITE_KEYS)

_UNSET = object()

def _contains_proxy(*values: Any) -> bool:
    for _proxy in _walk_proxies(values):
        return True
    return False

def _validate_name(name: str) -> tuple[str, str]:
    if not isinstance(name, str):
        raise TypeError(f"op name must be a str, got {type(name)!r}")
    parts = name.split("::")
    if len(parts) != 2 or not parts[0] or not parts[1]:
        raise ValueError(
            f'expected a qualified op name like "mylib::my_op", got {name!r}'
        )
    namespace, opname = parts
    if not all(part.isidentifier() for part in (namespace, opname)):
        raise ValueError(
            f"namespace and op name must be identifiers, got {name!r}"
        )
    return namespace, opname

def _normalize_device_types(device_types: Any) -> list[str] | None:
    """Normalize ``device_types`` into registry keys.

    ``device_types`` is omitted); strings and iterables of strings map onto
    per-device entries keyed by ``"cpu"``/``"cuda"``/...
    """

    if device_types is None:
        return None
    items = [device_types] if isinstance(device_types, str) else list(device_types)
    keys: list[str] = []
    for item in items:
        if not isinstance(item, str):
            raise TypeError(
                f"device types must be strings, got {item!r}; use None to "
                "declare one device-agnostic implementation"
            )
        key = item.lower()
        if key not in keys:
            keys.append(key)
    return keys

def _bridge_slot_key(device_types: Any) -> Any:
    """

    ``None``/empty iterables/composite spellings select the device-agnostic
    slot (``None`` key); concrete devices lowercase into their own slot.
    Used by the top-level registrations and ``Library.impl``.
    """

    if device_types is None:
        return None
    if isinstance(device_types, str):
        return None if device_type_is_composite(device_types) else device_types.lower()
    keys = list(device_types)
    if not keys:
        return None
    if len(keys) == 1:
        return _bridge_slot_key(keys[0])
    return [_bridge_slot_key(k) for k in keys]

def device_type_is_composite(device_type: str) -> bool:
    return device_type in _COMPOSITE_KEYS

def _validate_mutates_args(mutates_args: Any) -> tuple[str, ...]:
    if isinstance(mutates_args, str) or not isinstance(mutates_args, Iterable):
        raise TypeError(
            f"mutates_args must be an iterable of argument names, got "
            f"{mutates_args!r}"
        )
    mutated = tuple(mutates_args)
    if any(not isinstance(item, str) for item in mutated):
        raise TypeError(
            f"mutates_args entries must be strings, got {mutated!r}"
        )
    return mutated

def _validate_schema(schema: Any) -> str | None:
    """

    TensorPlay keeps them verbatim for introspection/opcheck, so only the
    qualified-name head is checked.
    """

    if schema is None:
        return None
    if not isinstance(schema, str):
        raise TypeError(f"schema must be a str or None, got {type(schema)!r}")
    signature = schema.split("(", 1)[0].strip()
    if signature:
        _validate_name(signature)
    elif not schema.lstrip().startswith("("):
        # Bare-fragment spelling: custom_op carries the qualified name in
        # its own argument, so the schema is just the signature.  A string
        # that is neither a qualified head nor a signature is a typo.
        raise ValueError(
            f'expected a schema like "(Tensor x) -> Tensor" or '
            f'"mylib::my_op(Tensor x) -> Tensor", got {schema!r}'
        )
    return schema

class CustomOpDef:
    """

    Instances are callable.  Calling with symbolic (tracer proxy) arguments
    records one opaque graph node; calling with real tensors dispatches to
    the kernel registered for the first tensor argument's device, wrapped
    in the registered autograd formula when gradients are requested.
    """

    def __init__(
        self,
        name: str,
        *,
        mutates_args: Sequence[str] = (),
        device_types: Any = None,
        schema: str | None = None,
        is_triton_op: bool = False,
        is_tile_lang_op: bool = False,
    ) -> None:
        self._namespace, self._opname = _validate_name(name)
        self._name = f"{self._namespace}::{self._opname}"
        self._mutates_args = frozenset(_validate_mutates_args(mutates_args))
        self._device_keys = _normalize_device_types(device_types)
        self._schema = _validate_schema(schema)
        self._is_triton_op = bool(is_triton_op)
        self._is_tile_lang_op = bool(is_tile_lang_op)
        # ``None`` key holds the device-agnostic kernel (device_types=None).
        self._kernels: dict[str | None, Callable[..., Any]] = {}
        self._disabled_kernels: set[str] = set()
        # Memoized _kernel_for resolutions, keyed by device key.  Every
        # mutation of _kernels/_disabled_kernels clears this (registration
        # and enable/disable are rare; the memo turns the per-call
        # resolution into a single dict hit).
        self._kernel_cache: dict[str | None, Callable[..., Any]] = {}
        self._fake_fn: Callable[..., Any] | None = None
        self._backward: Callable[..., Any] | None = None
        self._setup_context: Callable[..., Any] | None = None
        self._vmap_fn: Callable[..., Any] | None = None
        self._autocast_rules: dict[str, Any] = {}
        self._autograd_cls: type | None = None
        self._autograd_apply: Callable[..., Any] | None = None

    def _install_default_kernel(self, fn: Callable[..., Any]) -> None:
        """Use ``fn`` as the initial kernel (the ``@custom_op`` body).

        the advertised ``device_types`` — every device when omitted.
        """

        if not callable(fn):
            raise TypeError(f"operator body must be callable, got {type(fn)!r}")
        if self._device_keys is None:
            self._kernels[None] = fn
        else:
            for key in self._device_keys:
                self._kernels[key] = fn
        self._kernel_cache.clear()
        self._mirror_native(self._device_keys, fn)

    # -- introspection -----------------------------------------------------

    @property
    def name(self) -> str:
        return self._name

    @property
    def namespace(self) -> str:
        return self._namespace

    @property
    def opname(self) -> str:
        return self._opname

    @property
    def mutates_args(self) -> frozenset[str]:
        return self._mutates_args

    @property
    def is_triton_op(self) -> bool:
        return self._is_triton_op

    @property
    def is_tile_lang_op(self) -> bool:
        return self._is_tile_lang_op

    @property
    def schema(self) -> str | None:
        return self._schema

    # Readable node names once this object becomes a graph target
    # (the graph target formatter resolves ``target.__name__`` first).
    @property
    def __name__(self) -> str:  # type: ignore[override]
        return self._name

    def __repr__(self) -> str:
        kind = (
            "triton_op"
            if self._is_triton_op
            else "tile_lang_op"
            if self._is_tile_lang_op
            else "custom_op"
        )
        return f"<{kind} {self._name}>"

    # -- registration API --------------------------------------------------

    def register_kernel(
        self, device_types: Any = None, fn: Callable[..., Any] | None = None, /
    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
        """Register the implementation for one or more device types.

        ``device_types=None`` registers a single device-agnostic kernel;
        otherwise pass a device string (``"cpu"``, ``"cuda"``, ...) or an
        iterable of them.  Usable directly
        (``op.register_kernel("cpu", my_fn)``) or as a decorator
        (``@op.register_kernel("cpu")``).  Re-registering the same device
        allows hot swaps for interactive sessions and tests).

        CPU/CUDA kernels are additionally mirrored into the native p10
        dispatcher under this operator's qualified name, so native code and
        :meth:`run_native` can invoke them through the real dispatch path.
        """

        keys = _normalize_device_types(device_types)

        def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
            if not callable(f):
                raise TypeError(f"kernel must be callable, got {type(f)!r}")
            with _LOCK:
                if keys is None:
                    self._kernels[None] = f
                else:
                    for key in keys:
                        self._kernels[key] = f
                self._kernel_cache.clear()
            self._mirror_native(keys, f)
            return f

        return decorator(fn) if fn is not None else decorator

    def _mirror_native(
        self, keys: list[str] | None, fn: Callable[..., Any]
    ) -> None:
        """Push a kernel into the native dispatcher bridge (best effort).

        Tensor-in/tensor-out kernels on CPU/CUDA become callable from native
        code via ``Dispatcher::findHandle("ns::op")``; other signatures keep
        working through the Python dispatch path only.
        """

        bridge = getattr(tensorplay._C, "_register_python_op_kernel", None)
        if bridge is None:
            return
        slots = ["default"] if keys is None else [k for k in keys if k in ("cpu", "cuda")]
        for slot in slots:
            try:
                bridge(self._name, slot, fn)
            except Exception:  # noqa: BLE001 - this is an optimization
                return

    def register_fake(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """Register a meta/fake kernel computing output metadata.

        The fake kernel receives the same arguments but must not allocate
        real tensor data; it returns tensors describing shape/dtype/device
        (``tensorplay.empty_like`` style factories without data).  It is
        exercised by :func:`opcheck`'s ``test_faketensor``; capturing
        compilers never execute either kernel version during tracing.
        """

        if not callable(fn):
            raise TypeError(f"fake kernel must be callable, got {type(fn)!r}")
        self._fake_fn = fn
        return fn

    def register_autograd(
        self,
        backward: Callable[..., Any],
        /,
        *,
        setup_context: Callable[..., Any] | None = None,
    ) -> None:
        """

        ``backward(ctx, *grad_outputs)`` receives the saved context.  The
        ``setup_context`` callback may save tensors via
        ``ctx.save_for_backward``.  When no
        ``setup_context`` is given the context stays empty, so backward must
        derive its result purely from ``grad_outputs`` (or close over module
        state).
        """

        if not callable(backward):
            raise TypeError(f"backward must be callable, got {type(backward)!r}")
        if setup_context is not None and not callable(setup_context):
            raise TypeError(
                f"setup_context must be callable or None, got {setup_context!r}"
            )
        self._backward = backward
        self._setup_context = setup_context
        self._autograd_cls = self._build_autograd_class()
        self._autograd_apply = self._autograd_cls.apply

    def register_vmap(self, fn: Callable[..., Any]) -> Callable[..., Any]:
        """

        The engine does not batch yet; the registration is stored and surfaced
        through :func:`get_kernel`-style introspection.
        """

        if not callable(fn):
            raise TypeError(f"vmap function must be callable, got {type(fn)!r}")
        self._vmap_fn = fn
        # Rebuild so the generated autograd class exposes the formula the
        # same way a hand-written Function subclass would.
        if self._backward is not None:
            self._autograd_cls = self._build_autograd_class()
            self._autograd_apply = self._autograd_cls.apply
        return fn

    def register_autocast(
        self, device_type: str, cast_inputs: Any
    ) -> None:
        """

        When autocast is enabled for ``device_type``, floating-point tensor
        arguments are cast to ``cast_inputs`` before any kernel runs — the
        """

        if not isinstance(device_type, str):
            raise TypeError(
                f"device_type must be a str, got {type(device_type)!r}"
            )
        if not (
            hasattr(cast_inputs, "is_floating_point")
            and cast_inputs.is_floating_point
        ):
            raise TypeError(
                f"cast_inputs must be a floating-point dtype, got "
                f"{cast_inputs!r}"
            )
        self._autocast_rules[device_type.lower()] = cast_inputs

    @contextlib.contextmanager
    def set_kernel_enabled(self, device_type: str, enabled: bool = True):
        """

        context is active the concrete kernel for ``device_type`` is skipped
        and dispatch falls back to the device-agnostic kernel (if any).
        Disabling an already-disabled (or enabling an already-enabled)
        kernel warns and is otherwise a no-op; the original state is always
        restored on exit.
        """

        if isinstance(device_type, str):
            key = device_type.lower()
        else:
            key = device_type
        originally_disabled = key in self._disabled_kernels
        has_own_kernel = key in self._kernels
        action = "enable" if enabled else "disable"
        if not has_own_kernel and None not in self._kernels:
            warnings.warn(
                f"Attempted to {action} kernel for {key!r} but no kernel was "
                "registered for this device type.",
                stacklevel=2,
            )
        if not enabled:
            if originally_disabled:
                warnings.warn(
                    f"Attempted to disable kernel for {key!r} but it was "
                    "already disabled.",
                    stacklevel=2,
                )
            else:
                self._disabled_kernels.add(key)
                self._kernel_cache.clear()
        else:  # enable the kernel
            if not originally_disabled:
                warnings.warn(
                    f"Attempted to enable kernel for {key!r} but it was "
                    "already enabled.",
                    stacklevel=2,
                )
            else:
                self._disabled_kernels.remove(key)
                self._kernel_cache.clear()
        try:
            yield
        finally:
            # restore original state
            if originally_disabled:
                self._disabled_kernels.add(key)
            else:
                self._disabled_kernels.discard(key)
            self._kernel_cache.clear()

    def _build_autograd_class(self) -> type:
        global _CONTEXT_CLS
        if _CONTEXT_CLS is None:
            from tensorplay.autograd.function import _Context
            _CONTEXT_CLS = _Context
        op_def = self

        # Legacy forward shape (ctx first): the native apply helper then
        # makes ONE Python call per op instead of separate forward and
        # setup-context callbacks, and classes whose operator registers no
        # setup_context skip that work entirely.
        class _CustomOpAutograd(tensorplay.autograd.Function):
            @staticmethod
            def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
                output = op_def._run_kernel(args, kwargs)
                if op_def._setup_context is not None:
                    op_def._setup_context(ctx, args, output)
                return output

            @staticmethod
            def backward(ctx: Any, *grad_outputs: Any) -> Any:
                return op_def._backward(ctx, *grad_outputs)

        if op_def._vmap_fn is not None:
            _vmap_fn = op_def._vmap_fn

            @staticmethod
            def vmap(info: Any, in_dims: Any, *args: Any) -> Any:
                return _vmap_fn(info, in_dims, *args)

            _CustomOpAutograd.vmap = staticmethod(vmap)  # type: ignore[method-assign]

        _CustomOpAutograd.__qualname__ = f"_CustomOpAutograd[{self._name}]"
        return _CustomOpAutograd

    # -- dispatch ----------------------------------------------------------

    def _kernel_for(
        self, args: tuple[Any, ...], key: Any = _UNSET
    ) -> Callable[..., Any]:
        # Lock-free read path: dict/set membership tests are atomic under
        # the GIL and registration only ever swaps callables wholesale.
        if key is _UNSET:
            key = _first_device_key(args)
        # Hot path: one dict hit for a previously resolved (key, state)
        # combination.  _disabled_kernels membership dominates the miss
        # path only when a kernel was toggled; both mutation sites clear
        # the memo.
        cache = self._kernel_cache
        fn = cache.get(key)
        if fn is not None:
            return fn
        kernels = self._kernels
        if key is not None:
            fn = kernels.get(key)
            if fn is not None and key not in self._disabled_kernels:
                cache[key] = fn
                return fn
            # Disabled/shadowed concrete kernels fall back to the composite
        fn = kernels.get(None)
        if fn is not None:
            if key is not None and key not in self._disabled_kernels:
                cache[key] = fn
            return fn
        with _LOCK:  # cold error path only
            registered = sorted(str(item) for item in kernels)
        where = f" for device {key!r}" if key is not None else ""
        hint = f"; registered devices: {registered}" if registered else ""
        raise NotImplementedError(
            f"{self._name} has no kernel{where}{hint}. Register one via "
            f"{self._name}.register_kernel(...)"
        )

    def _run_kernel(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
        return self._kernel_for(args)(*args, **kwargs)

    def run_native(
        self, inputs: Sequence[Any], *, device_type: str | None = None
    ) -> list[tensorplay.Tensor]:
        """Invoke this operator through the native p10 dispatcher.

        Exercises ``Dispatcher::findHandle`` plus the real kernel table —
        the same path C++ code takes when calling the operator by name.
        Requires a tensor-only signature (the canonical unboxed convention
        is tensors-in/tensors-out); kernels that were never mirrored (e.g.
        non-tensor arguments) raise ``NotImplementedError``.
        """

        bridge = getattr(tensorplay._C, "_call_native_op", None)
        if bridge is None:
            raise RuntimeError(
                "the native custom-op bridge is unavailable in this build"
            )
        return bridge(self._name, list(inputs), device_type)

    def _autograd_call(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
        """Registered-op call through the one-crossing native apply.

        The native side invokes the generated forward (which runs the user
        setup_context inside the disabled-grad block), wires the context and
        attaches the graph node.  Returns None when no input needs grad (or
        grads are disabled); the caller then runs the kernel directly
        instead of building and discarding an empty graph node.
        """
        out = _CUSTOM_OP_AUTOGRAD_APPLY(
            _CONTEXT_CLS, self._autograd_cls.forward, self._backward,
            self._autograd_cls._node_name, args, kwargs)
        if out is None:
            return self._run_profiled(args, kwargs, None)
        return out

    def _eager_call(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
        """Dispatch real tensors with full eager semantics (no capture).

        Shared by :meth:`__call__` and the native re-entry below so compiled
        graphs keep device dispatch AND ``register_autograd`` behavior.
        Hot path: one device-key scan, one dict lookup, one call.  The
        autocast rule table stays empty for the common un-registered case,
        so its branch is fully skipped there.
        """

        if self._autograd_cls is not None:
            # The autocast cast preserves requires_grad, so it can run on
            # the raw arguments before the native gate below.
            if self._autocast_rules:
                key = _first_device_key(args)
                if key is not None:
                    rule = self._autocast_rules.get(key)
                    if rule is not None and _autocast_enabled(key):
                        args = tuple(_cast_if_floating(v, rule) for v in args)
                        if kwargs:
                            kwargs = {
                                k: _cast_if_floating(v, rule)
                                for k, v in kwargs.items()
                            }
            if _profiling_sessions:
                return self._profiled(self._autograd_call, args, kwargs)
            if _CUSTOM_OP_AUTOGRAD_APPLY is not None:
                return self._autograd_call(args, kwargs)
            # Stale-extension fallback: Python gate plus Function.apply.
            # Keyword arguments are scanned only when present: the unpacking
            # itself would cost more than the scan on the common path.
            if _is_grad_enabled() and (
                    _any_requires_grad(*args) if not kwargs
                    else _any_requires_grad(*args, **kwargs)):
                return self._autograd_apply(*args, **kwargs)
            return self._run_profiled(args, kwargs, None)
        if self._autocast_rules:
            key = _first_device_key(args)
            rule = self._autocast_rules.get(key) if key is not None else None
            if rule is not None and _autocast_enabled(key):
                args = tuple(_cast_if_floating(v, rule) for v in args)
                if kwargs:
                    kwargs = {
                        k: _cast_if_floating(v, rule) for k, v in kwargs.items()
                    }
            return self._run_profiled(args, kwargs, key)
        return self._run_profiled(args, kwargs, None)

    # -- profiler span emission --------------------------------------------

    def _profiled(self, call, args, kwargs):
        """Run ``call`` inside one user-annotation span (autograd path)."""
        begin = tensorplay._C._profiler_user_begin
        end = tensorplay._C._profiler_user_end
        begin(self._name)
        try:
            return call(*args, **kwargs)
        finally:
            end()

    def _run_profiled(
        self, args: tuple[Any, ...], kwargs: dict[str, Any], key: Any
    ) -> Any:
        """Kernel invocation, emitting one op span when a session is live.

        The session flag is a module-global load; outside profiling this is
        the plain ``_kernel_for`` call with no extra work.  Inner calls the
        kernel itself fires (e.g. a composite body) nest under this span
        like ordinary operator records do.
        """
        if _profiling_sessions:
            begin = tensorplay._C._profiler_user_begin
            end = tensorplay._C._profiler_user_end
            begin(self._name)
            try:
                return (self._kernel_for(args) if key is None else self._kernel_for(args, key))(
                    *args, **kwargs
                )
            finally:
                end()
        if key is None:
            return self._kernel_for(args)(*args, **kwargs)
        return self._kernel_for(args, key)(*args, **kwargs)

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        # A Proxy can only exist while a Tracer.trace() is live on this
        # thread (the graph tracing depth), so the proxy walk is skipped
        # entirely outside capture — identical to the generated functional
        # wrappers' hot path in tensorplay/functional.py.
        if _capturing():
            captured = _capture_call(self, args, kwargs)
            if captured is not None:
                return captured
        mode = _recording_mode()
        if mode is not None:
            # A library operation is a node in a graph just as an operation
            # reached through the dispatcher is, and the operation that reads
            # randomness takes its position from a pair of tensors rather than
            # from a proxy -- so the walk above finds nothing to record it
            # from, and the values it produced would be held by name in the
            # graph instead of being tied to the node that produced them.
            #
            # The recording is done here rather than by handing the call to a
            # mode on the stack: a mode that stands values in for real ones
            # answers a call by making it again, and that call is this one.
            #
            # The body runs with the recording mode off the stack, as an
            # operator's kernel runs inside a mode's handler: the operations
            # it makes belong to this node, and recording or decomposing them
            # again would put them in the graph twice -- or, for a body that
            # calls the operation this one decomposes from, never finish.
            with _modes_popped_through(mode):
                out = self._eager_call(args, kwargs)
            mode.tracer.record(self, args, kwargs, out)
            return out
        return self._eager_call(args, kwargs)

def _recording_mode() -> Any:
    """The innermost mode that records into a graph being traced, if any.

    A mode on the stack is not by itself a recorder: a mode that stands values
    in for real ones answers an operation by running it against the stand-ins,
    and the tracer of a mode that records is the one that knows which graph is
    being written.  Only a mode that carries such a tracer counts.
    """

    # An empty mode stack can carry no tracer; the length probe is one
    # native call, far cheaper than materializing the stack.
    if _len_dispatch_modes is not None and _len_dispatch_modes() == 0:
        return None

    from .utils._dispatch import _get_current_dispatch_mode_stack

    for mode in reversed(_get_current_dispatch_mode_stack()):
        tracer = getattr(mode, "tracer", None)
        if tracer is not None and hasattr(tracer, "record"):
            return mode
    return None

@contextlib.contextmanager
def _modes_popped_through(mode: Any):
    """Take ``mode`` and every mode above it off the stack for the block."""

    popped = []
    try:
        while True:
            top = tensorplay._C._pop_dispatch_mode()
            popped.append(top)
            if top is mode:
                break
        yield
    finally:
        for top in reversed(popped):
            tensorplay._C._push_dispatch_mode(top)

def _first_device_key(values: tuple[Any, ...]) -> str | None:
    if _fast_device_key is not None:
        return _fast_device_key(*values)
    for value in values:
        if isinstance(value, tensorplay.Tensor):
            device = value.device
            device_type = getattr(device, "type", None)
            if isinstance(device_type, str):
                return device_type
            return "cuda" if device.is_cuda() else "cpu"
    return None

def _autocast_enabled(device_key: str) -> bool:
    fn = getattr(tensorplay, "is_autocast_enabled", None)
    if fn is None:
        return False
    try:
        return bool(fn(device_key))
    except Exception:  # noqa: BLE001 - autocast is best effort
        return False

def _cast_if_floating(value: Any, dtype: Any) -> Any:
    if isinstance(value, tensorplay.Tensor):
        value_dtype = value.dtype
        if value_dtype.is_floating_point and value_dtype != dtype:
            return value.to(dtype)
    return value

def get_op(op_name: str) -> CustomOpDef:
    """Return the previously defined operator named ``"ns::op"``."""

    if not isinstance(op_name, str):
        raise TypeError(f"op name must be a str, got {type(op_name)!r}")
    with _LOCK:
        op_def = _OP_REGISTRY.get(op_name)
    if op_def is None:
        raise RuntimeError(f"unknown operator {op_name!r}")
    return op_def

def has_op(op_name: str) -> bool:
    with _LOCK:
        return op_name in _OP_REGISTRY

def _resolve_op(op: str | CustomOpDef) -> CustomOpDef:
    if isinstance(op, CustomOpDef):
        return op
    if isinstance(op, str):
        return get_op(op)
    raise TypeError(
        f"expected an operator name or CustomOpDef, got {type(op)!r}"
    )

def _register_op_def(
    op_def: CustomOpDef,
) -> CustomOpDef:
    with _LOCK:
        previous = _OP_REGISTRY.get(op_def.name)
        if previous is not None:
            raise RuntimeError(
                f"operator {op_def.name!r} is already defined; use "
                f"get_op({op_def.name!r}) or a fresh namespace"
            )
        _OP_REGISTRY[op_def.name] = op_def
    return op_def

def _define_or_get(qualname: str, schema: str | None) -> CustomOpDef:
    namespace, _opname = _validate_name(qualname)
    del namespace
    _validate_schema(schema)
    with _LOCK:
        op_def = _OP_REGISTRY.get(qualname)
    if op_def is None:
        op_def = _register_op_def(CustomOpDef(qualname))
    elif schema is not None:
        op_def._schema = _validate_schema(schema)
    return op_def

[docs]
def 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]:
    """

    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).

    Args:
        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
            :meth:`CustomOpDef.register_kernel`.
        schema: Optional schema string kept for introspection and
            :func:`opcheck` (TensorPlay models no schema grammar).

    Returns:
        A decorator producing a callable :class:`CustomOpDef`.
    """

    _validate_name(name)
    _validate_mutates_args(mutates_args)
    _normalize_device_types(device_types)
    _validate_schema(schema)

    def decorator(f: Callable[..., Any]) -> CustomOpDef:
        op_def = CustomOpDef(
            name,
            mutates_args=mutates_args,
            device_types=device_types,
            schema=schema,
        )
        _register_op_def(op_def)
        op_def._install_default_kernel(f)
        return op_def

    return decorator(fn) if fn is not None else decorator

[docs]
def triton_op(
    name: str,
    fn: Callable[..., Any] | None = None,
    /,
    *,
    mutates_args: Sequence[str] = (),
    device_types: Any = None,
    schema: str | None = None,
) -> Callable[[Callable[..., Any]], CustomOpDef]:
    """

    The registered kernel(s) must launch their Triton kernels through
    :func:`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
    """

    _validate_name(name)
    _validate_mutates_args(mutates_args)
    _normalize_device_types(device_types)
    _validate_schema(schema)

    def decorator(f: Callable[..., Any]) -> CustomOpDef:
        op_def = CustomOpDef(
            name,
            mutates_args=mutates_args,
            device_types=device_types,
            schema=schema,
            is_triton_op=True,
        )
        _register_op_def(op_def)
        op_def._install_default_kernel(f)
        return op_def

    return decorator(fn) if fn is not None else decorator

def tile_lang_op(
    name: str,
    fn: Callable[..., Any] | None = None,
    /,
    *,
    mutates_args: Sequence[str] = (),
    device_types: Any = None,
    schema: str | None = None,
) -> Callable[[Callable[..., Any]], CustomOpDef]:
    """Define a TileLang-backed operator (Triton contract, tile-lang kernels).

    TileLang (https://github.com/tile-ai/tilelang) compiles ``@T.prim_func``
    DSL programs into highly-tuned CUDA/Metal/CPU kernels; grid and thread
    configuration live inside the prim func's ``with T.Kernel(...)`` block,
    so unlike Triton there is no grid indexing at the launch site — a
    compiled ``JITKernel`` is called directly with tensors.

    The registered kernel body must launch its TileLang kernels through
    :func:`wrap_tilelang`.  Under ``tensorplay.compile`` the operator is
    captured as a single opaque fusion-barrier node whose body never runs
    during tracing, exactly like :func:`triton_op`.
    """

    _validate_name(name)
    _validate_mutates_args(mutates_args)
    _normalize_device_types(device_types)
    _validate_schema(schema)

    def decorator(f: Callable[..., Any]) -> CustomOpDef:
        op_def = CustomOpDef(
            name,
            mutates_args=mutates_args,
            device_types=device_types,
            schema=schema,
            is_tile_lang_op=True,
        )
        _register_op_def(op_def)
        op_def._install_default_kernel(f)
        return op_def

    return decorator(fn) if fn is not None else decorator

def _native_invoke(op_name: str, *tensors: Any) -> Any:
    """Re-entry point for compiled native graphs (Stax ``custom_op`` nodes).

    The C++ executor installed by the bindings calls this with the
    operator's qualified name and its tensor inputs; routing through the
    :class:`CustomOpDef` keeps device dispatch and autograd identical to
    eager execution instead of bypassing them with a raw kernel.
    """

    return get_op(op_name)._eager_call(tuple(tensors), {})

class TritonKernelWrapper:
    """Grid-indexable passthrough around a ``@triton.jit`` kernel.

    At eager time ``wrapped[grid](...)`` simply launches the kernel.  If a
    launch is ever captured symbolically (proxy arguments), it raises:
    Triton launches are never part of the canonical graph — they live
    inside :func:`triton_op` bodies, which the compiler captures as one
    opaque node instead.
    """

    __slots__ = ("kernel",)

    def __init__(self, kernel: Any) -> None:
        self.kernel = kernel

    def __getitem__(self, grid: Any) -> Callable[..., Any]:
        kernel = self.kernel

        def launcher(*args: Any, **kwargs: Any) -> Any:
            if _contains_proxy(grid, args, kwargs):
                raise GraphCaptureError(
                    "a raw Triton launch cannot be captured inside "
                    "tensorplay.compile; define the launch inside "
                    "tensorplay.library.triton_op and call that operator "
                    "instead"
                )
            return kernel[grid](*args, **kwargs)

        launcher.__name__ = getattr(kernel, "__name__", "triton_kernel")
        return launcher

[docs]
def wrap_triton(kernel: Any) -> TritonKernelWrapper:
    """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.
    """

    if isinstance(kernel, TritonKernelWrapper):
        return kernel
    is_jit_function = False
    try:
        from triton.runtime.jit import JITFunction

        is_jit_function = isinstance(kernel, JITFunction)
    except ImportError:
        is_jit_function = False
    if not is_jit_function and not hasattr(kernel, "__getitem__"):
        raise TypeError(
            "wrap_triton expects a @triton.jit kernel (JITFunction) or a "
            f"grid-indexable launcher, got {type(kernel)!r}"
        )
    return TritonKernelWrapper(kernel)

class TileLangKernelWrapper:
    """Passthrough around a compiled TileLang kernel.

    TileLang launch sites carry no grid (``with T.Kernel(...)`` lives in
    the prim func), so unlike :class:`TritonKernelWrapper` this wraps a
    directly-callable object: a ``tilelang.jit.kernel.JITKernel``, the
    ``JITImpl`` produced by ``@tilelang.jit`` (lazy mode: calling the
    factory with shape/constexpr arguments yields the compiled kernel),
    or any duck-typed adapter.  Eager calls forward untouched; a symbolic
    (proxy) launch raises, steering raw launches behind a
    :func:`tile_lang_op` boundary.
    """

    __slots__ = ("kernel",)

    def __init__(self, kernel: Any) -> None:
        self.kernel = kernel

    def compile(self, *args: Any, **kwargs: Any) -> "TileLangKernelWrapper":
        """Bind a lazy-mode ``@tilelang.jit`` factory into a ready kernel."""

        bound = self.kernel(*args, **kwargs)
        return TileLangKernelWrapper(bound)

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if _contains_proxy(args, kwargs):
            raise GraphCaptureError(
                "a raw TileLang launch cannot be captured inside "
                "tensorplay.compile; define the launch inside "
                "tensorplay.library.tile_lang_op and call that operator "
                "instead"
            )
        return self.kernel(*args, **kwargs)

def _looks_like_tilelang_kernel(kernel: Any) -> bool:
    """Duck-typing for TileLang objects without importing tilelang.

    Recognizes ``tilelang.jit.kernel.JITKernel`` (compiled, exposes
    factory, exposes ``get_tir``/``out_idx``) and plain adapters that opt
    in via a ``_tilelang_kernel`` marker attribute.
    """

    if hasattr(kernel, "_tilelang_kernel"):
        return bool(getattr(kernel, "_tilelang_kernel"))
    return any(
        hasattr(kernel, attr) for attr in ("adapter", "torch_function", "get_tir")
    )

def wrap_tilelang(kernel: Any) -> TileLangKernelWrapper:
    """Mark a TileLang kernel as launchable from within a ``tile_lang_op``.

    Accepts a compiled ``tilelang.jit.kernel.JITKernel``, the lazy-mode
    ``JITImpl`` factory returned by ``@tilelang.jit``, or any duck-typed
    callables are rejected so typos surface early.  Idempotent on
    wrappers.
    """

    if isinstance(kernel, TileLangKernelWrapper):
        return kernel
    recognized = False
    try:
        from tilelang.jit.kernel import JITKernel as _JITKernel

        recognized = isinstance(kernel, _JITKernel)
    except ImportError:
        recognized = False
    if not recognized:
        try:
            from tilelang.jit import JITImpl as _JITImpl

            recognized = isinstance(kernel, _JITImpl)
        except ImportError:
            recognized = False
    if not recognized:
        recognized = _looks_like_tilelang_kernel(kernel)
    if not recognized and callable(kernel):
        raise TypeError(
            "wrap_tilelang expects a tilelang JITKernel/JITImpl (see "
            "https://github.com/tile-ai/tilelang) or a duck-typed adapter "
            f"{type(kernel)!r}; plain callables are rejected — set "
            f"'_tilelang_kernel = True' on custom launchers to opt in"
        )
    if not callable(kernel):
        raise TypeError(
            "wrap_tilelang expects a callable kernel, got "
            f"{type(kernel)!r}"
        )
    return TileLangKernelWrapper(kernel)

[docs]
def register_kernel(
    op: str | CustomOpDef,
    device_types: Any = None,
    func: Callable[..., Any] | None = None,
    /,
    *,
    lib: Any = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """

    Accepts a :class:`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.
    """

    del lib
    op_def = _resolve_op(op)
    key: Any = _bridge_slot_key(device_types)
    return op_def.register_kernel(key) if func is None else op_def.register_kernel(key)(func)

[docs]
def register_fake(
    op: str | CustomOpDef,
    func: Callable[..., Any] | None = None,
    /,
    *,
    lib: Any = None,
    allow_override: bool = True,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:

    del lib, allow_override
    op_def = _resolve_op(op)
    return op_def.register_fake(func) if func is not None else op_def.register_fake

[docs]
def register_autograd(
    op: str | CustomOpDef,
    backward: Callable[..., Any],
    /,
    *,
    setup_context: Callable[..., Any] | None = None,
    lib: Any = None,
) -> None:

    del lib
    op_def = _resolve_op(op)
    op_def.register_autograd(backward, setup_context=setup_context)

[docs]
def register_vmap(
    op: str | CustomOpDef,
    func: Callable[..., Any] | None = None,
    /,
    *,
    lib: Any = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:

    del lib
    op_def = _resolve_op(op)
    return op_def.register_vmap(func) if func is not None else op_def.register_vmap

[docs]
def register_autocast(
    op: str | CustomOpDef,
    device_type: str,
    cast_inputs: Any,
    /,
    *,
    lib: Any = None,
) -> None:

    del lib
    _resolve_op(op).register_autocast(device_type, cast_inputs)

[docs]
def define(
    qualname: str, schema: str | None = None, *, lib: Any = None, tags: Any = ()
) -> None:
    """

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

    del lib, tags
    if isinstance(qualname, str) and "(" in qualname:
        if schema is None:
            schema = qualname
        qualname = qualname.split("(", 1)[0].strip()
    _define_or_get(qualname, schema)

[docs]
def impl(
    qualname: str,
    types: Any,
    func: Callable[..., Any] | None = None,
    /,
    *,
    lib: Any = None,
) -> Callable[..., Any] | Callable[[Callable[..., Any]], Callable[..., Any]]:
    """

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

    del lib
    return register_kernel(qualname, types, func)

[docs]
def impl_abstract(
    qualname: str,
    func: Callable[..., Any] | None = None,
    /,
    *,
    lib: Any = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:

    return register_fake(qualname, func, lib=lib)

[docs]
def get_kernel(
    op: str | CustomOpDef, dispatch_key: str
) -> Callable[..., Any]:
    """

    ``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
    :meth:`CustomOpDef.set_kernel_enabled` visibility).
    """

    op_def = _resolve_op(op)
    if not isinstance(dispatch_key, str):
        raise TypeError(
            f"dispatch_key must be a str, got {type(dispatch_key)!r}"
        )
    lowered = dispatch_key.lower()
    key: Any = (
        None
        if lowered in _COMPOSITE_KEYS_LOWERED or lowered in ("default", "composite")
        else lowered
    )
    # does NOT satisfy a concrete-device query (use "default" for that).
    fn = op_def._kernels.get(key)
    if fn is not None and (
        key is None or key not in op_def._disabled_kernels
    ):
        return fn
    raise LookupError(
        f"No kernel registered for {op_def.name} with dispatch key "
        f"{dispatch_key!r}"
    )

# ---------------------------------------------------------------------------
# infer_schema
# ---------------------------------------------------------------------------

_SCHEMA_PRIMITIVES: dict[Any, str] = {int: "SymInt", float: "float", bool: "bool", str: "str"}

def _annotation_to_schema_str(annotation: Any) -> str:
    """Best-effort translation of Python annotations to schema atoms.

    Unannotated parameters default to ``Tensor`` (the overwhelming custom-op
    case); unsupported annotations fall back to ``Tensor`` rather than
    failing, since TensorPlay enforces no grammar.
    """

    if annotation is None or annotation is type(None):
        return "Tensor"
    if annotation is tensorplay.Tensor or annotation == "Tensor":
        return "Tensor"
    if annotation in _SCHEMA_PRIMITIVES:
        return _SCHEMA_PRIMITIVES[annotation]
    origin = typing.get_origin(annotation)
    if origin is typing.Union or (
        hasattr(typing, "UnionType") and origin is typing.UnionType
    ):
        args = [a for a in typing.get_args(annotation) if a is not type(None)]
        if len(args) == 1:
            return _annotation_to_schema_str(args[0]) + "?"
        return "(" + ", ".join(_annotation_to_schema_str(a) for a in args) + ")"
    if origin in (list, Sequence, tuple):
        args = typing.get_args(annotation)
        if origin is tuple and len(args) >= 2:
            inner = ", ".join(_annotation_to_schema_str(a) for a in args)
            return f"({inner})"
        atom = _annotation_to_schema_str(args[0] if args else tensorplay.Tensor)
        return f"{atom}[]"
    text = str(annotation)
    if "Tensor" in text:
        return "Tensor"
    return "Tensor"

[docs]
def infer_schema(
    prototype_function: Callable[..., Any],
    /,
    *,
    mutates_args: Sequence[str],
    op_name: str | None = None,
    tags: Any = (),
) -> str:
    """

    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.
    """

    del tags
    mutated = frozenset(_validate_mutates_args(mutates_args))
    if not callable(prototype_function):
        raise TypeError(
            f"prototype_function must be callable, got "
            f"{type(prototype_function)!r}"
        )
    try:
        hints = typing.get_type_hints(prototype_function)
    except Exception:  # noqa: BLE001 - unresolvable hints fall back to raw
        hints = dict(getattr(prototype_function, "__annotations__", {}) or {})
    parameters = inspect.signature(prototype_function).parameters
    alias_letters = iter("abcdefghijklmnopqrstuvwxyz")
    parts: list[str] = []
    for pname, parameter in parameters.items():
        if parameter.kind in (
            inspect.Parameter.VAR_POSITIONAL,
            inspect.Parameter.VAR_KEYWORD,
        ):
            continue
        atom = _annotation_to_schema_str(hints.get(pname))
        if pname in mutated:
            atom = f"{atom}({next(alias_letters)}!)"
        parts.append(atom)
    return_atom = _annotation_to_schema_str(hints.get("return"))
    name = op_name if op_name is not None else getattr(
        prototype_function, "__name__", "op"
    )
    _validate_name(name) if "::" in name else None
    return f"{name}({', '.join(parts)}) -> {return_atom}"

# ---------------------------------------------------------------------------
# opcheck
# ---------------------------------------------------------------------------

_OPCHECK_DEFAULT_UTILS = (
    "test_schema",
    "test_autograd_registration",
    "test_faketensor",
    "test_aot_dispatch_dynamic",
)

def _flatten_tensors(value: Any) -> list[tensorplay.Tensor]:
    if isinstance(value, tensorplay.Tensor):
        return [value]
    if isinstance(value, (list, tuple)):
        flat: list[tensorplay.Tensor] = []
        for item in value:
            flat.extend(_flatten_tensors(item))
        return flat
    if isinstance(value, dict):
        flat = []
        for item in value.values():
            flat.extend(_flatten_tensors(item))
        return flat
    return []

def _bind_named_tensors(
    op_def: CustomOpDef, args: tuple[Any, ...], kwargs: dict[str, Any]
) -> dict[str, tensorplay.Tensor]:
    """Map argument names to tensor values via the kernel's signature."""

    kernel = next(iter(op_def._kernels.values()), None)
    named: dict[str, tensorplay.Tensor] = {}
    if kernel is not None:
        try:
            sig = inspect.signature(kernel)
            bound = sig.bind_partial(*args, **kwargs)
        except TypeError:
            bound = None
        if bound is not None:
            for name, value in bound.arguments.items():
                if isinstance(value, tensorplay.Tensor):
                    named[name] = value
                elif isinstance(value, (list, tuple)):
                    for item in value:
                        if isinstance(item, tensorplay.Tensor):
                            named[name] = item
                            break
            return named
    # Signature-less fallback: positional tensors get synthetic names.
    for index, value in enumerate(args):
        if isinstance(value, tensorplay.Tensor):
            named[f"arg{index}"] = value
    for name, value in kwargs.items():
        if isinstance(value, tensorplay.Tensor):
            named[name] = value
    return named

def _tensors_equal(a: tensorplay.Tensor, b: tensorplay.Tensor) -> bool:
    if a.shape != b.shape or a.dtype != b.dtype:
        return False
    try:
        return bool(tensorplay.equal(a, b))
    except Exception:  # noqa: BLE001 - exotic dtypes: fall back to identity
        return a.data_ptr() == b.data_ptr()

def _opcheck_test_schema(
    op_def: CustomOpDef,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> None:
    named_before = _bind_named_tensors(op_def, args, kwargs)
    snapshots = {n: t.clone() for n, t in named_before.items()}
    output = op_def(*args, **kwargs)
    # Undeclared mutation is the safety-critical direction: fusion treats
    # undeclared side effects as barriers, so catch silent in-place edits.
    for name, before in snapshots.items():
        if name in op_def.mutates_args:
            continue
        after = named_before[name]
        if not _tensors_equal(before, after):
            raise AssertionError(
                f"{op_def.name}: kernel mutated argument '{name}' which is "
                "not declared in mutates_args; declare it or stop mutating"
            )
    # Outputs must be fresh allocations, never aliases of an input storage
    input_ptrs = {t.data_ptr() for t in named_before.values()}
    for tensor in _flatten_tensors(output):
        if tensor.data_ptr() in input_ptrs:
            raise AssertionError(
                f"{op_def.name}: kernel returned a tensor aliasing an input; "
                "custom ops must return fresh outputs (mutation goes through "
                "mutates_args)"
            )

def _opcheck_test_faketensor(
    op_def: CustomOpDef,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> None:
    fake_fn = op_def._fake_fn
    if fake_fn is None:
        raise AssertionError(
            f"{op_def.name}: no fake kernel registered; register one with "
            "register_fake so compilers can propagate metadata"
        )
    real_out = _flatten_tensors(op_def(*args, **kwargs))
    fake_out = _flatten_tensors(fake_fn(*args, **kwargs))
    if len(real_out) != len(fake_out):
        raise AssertionError(
            f"{op_def.name}: fake kernel returned {len(fake_out)} tensor(s) "
            f"but the real kernel returned {len(real_out)}"
        )
    for i, (real, fake) in enumerate(zip(real_out, fake_out)):
        if tuple(real.shape) != tuple(fake.shape):
            raise AssertionError(
                f"{op_def.name}: output {i} shape mismatch "
                f"(real {tuple(real.shape)} vs fake {tuple(fake.shape)})"
            )
        if real.dtype != fake.dtype:
            raise AssertionError(
                f"{op_def.name}: output {i} dtype mismatch "
                f"(real {real.dtype} vs fake {fake.dtype})"
            )
        if str(real.device) != str(fake.device):
            raise AssertionError(
                f"{op_def.name}: output {i} device mismatch "
                f"(real {real.device} vs fake {fake.device})"
            )

def _opcheck_test_autograd_registration(
    op_def: CustomOpDef,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> None:
    kernel = next(iter(op_def._kernels.values()), None)
    if kernel is None:
        return
    try:
        parameters = [
            p
            for p in inspect.signature(kernel).parameters.values()
            if p.kind
            not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
        ]
        if len(parameters) < len(args):
            raise TypeError
    except (TypeError, ValueError):
        return  # cannot map inputs onto the kernel signature; skip probe

    new_args = list(args)
    new_kwargs = dict(kwargs)
    leaves: list[tuple[tensorplay.Tensor, tensorplay.Tensor]] = []
    for index, parameter in enumerate(parameters):
        source = args[index] if index < len(args) else new_kwargs.get(parameter.name)
        if (
            isinstance(source, tensorplay.Tensor)
            and source.dtype.is_floating_point
        ):
            leaf = source.detach().clone().requires_grad_(True)
            if index < len(args):
                new_args[index] = leaf
            else:
                new_kwargs[parameter.name] = leaf
            leaves.append((source, leaf))
    if not leaves:
        return  # nothing differentiable to probe

    outputs = op_def(*new_args, **new_kwargs)
    flat_outputs = _flatten_tensors(outputs)
    if not flat_outputs:
        return
    first = flat_outputs[0]
    if not first.requires_grad:
        raise AssertionError(
            f"{op_def.name}: output does not require grad although a "
            "floating input does; the kernel breaks the autograd graph"
        )
    first.backward(tensorplay.ones_like(first))
    for original, leaf in leaves:
        grad = leaf.grad
        if grad is None:
            raise AssertionError(
                f"{op_def.name}: no gradient reached input; register an "
                "autograd formula via register_autograd"
            )
        if tuple(grad.shape) != tuple(original.shape):
            raise AssertionError(
                f"{op_def.name}: gradient shape {tuple(grad.shape)} does not "
                f"match input shape {tuple(original.shape)}"
            )

def _make_fixed_arity_wrapper(
    fn: Callable[..., Any], nargs: int, kwargs: dict[str, Any]
) -> Callable[..., Any]:
    """Build a fixed-arity ``f(a0, ..., aN)`` shim (the tracer rejects
    varargs), forwarding into ``fn``."""
    if nargs == 0:
        return lambda: fn(**dict(kwargs))
    if nargs == 1:
        return lambda a0: fn(a0, **dict(kwargs))
    if nargs == 2:
        return lambda a0, a1: fn(a0, a1, **dict(kwargs))
    if nargs == 3:
        return lambda a0, a1, a2: fn(a0, a1, a2, **dict(kwargs))
    if nargs == 4:
        return lambda a0, a1, a2, a3: fn(a0, a1, a2, a3, **dict(kwargs))
    if nargs == 5:
        return lambda a0, a1, a2, a3, a4: fn(a0, a1, a2, a3, a4, **dict(kwargs))
    if nargs == 6:
        return lambda a0, a1, a2, a3, a4, a5: fn(
            a0, a1, a2, a3, a4, a5, **dict(kwargs)
        )
    raise ValueError(
        f"opcheck supports up to 6 sample tensor arguments, got {nargs}"
    )

def _opcheck_test_aot_dispatch_dynamic(
    op_def: CustomOpDef,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> None:
    from tensorplay.graph import Tracer

    tensor_positions = [
        i for i, v in enumerate(args) if isinstance(v, tensorplay.Tensor)
    ]
    if not tensor_positions:
        return  # nothing symbolic could flow; capture adds no information
    limit = min(len(tensor_positions), 6)
    samples = [args[i].clone() for i in tensor_positions[:limit]]
    wrapped = _make_fixed_arity_wrapper(op_def, len(samples), kwargs)
    param_names = [f"a{i}" for i in range(len(samples))]
    traced = Tracer().trace(
        wrapped, sample_inputs=dict(zip(param_names, samples))
    )
    eager_out = _flatten_tensors(op_def(*samples, **kwargs))
    compiled_out = _flatten_tensors(traced(*samples))
    if len(eager_out) != len(compiled_out):
        raise AssertionError(
            f"{op_def.name}: captured graph returned {len(compiled_out)} "
            f"tensor(s), eager returned {len(eager_out)}"
        )
    for i, (want, got) in enumerate(zip(eager_out, compiled_out)):
        if not bool(tensorplay.allclose(want, got)):
            raise AssertionError(
                f"{op_def.name}: compiled output {i} diverges from eager"
            )

_OPCHECK_TESTS = {
    "test_schema": _opcheck_test_schema,
    "test_faketensor": _opcheck_test_faketensor,
    "test_autograd_registration": _opcheck_test_autograd_registration,
    "test_aot_dispatch_dynamic": _opcheck_test_aot_dispatch_dynamic,
}

[docs]
def opcheck(
    op: str | CustomOpDef,
    args: tuple[Any, ...],
    kwargs: dict[str, Any] | None = None,
    *,
    test_utils: str | Sequence[str] = _OPCHECK_DEFAULT_UTILS,
    raise_exception: bool = True,
    atol: float | None = None,
    rtol: float | None = None,
) -> dict[str, str]:
    """

    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.
    """

    del atol, rtol  # accepted for signature compatibility; comparisons are exact
    op_def = _resolve_op(op)
    kwargs = dict(kwargs or {})
    if isinstance(test_utils, str):
        selected = (
            _OPCHECK_DEFAULT_UTILS if test_utils == "all" else (test_utils,)
        )
    else:
        selected = tuple(test_utils)
    unknown = [t for t in selected if t not in _OPCHECK_TESTS]
    if unknown:
        raise ValueError(
            f"unknown opcheck test_utils {unknown}; expected a subset of "
            f"{sorted(_OPCHECK_TESTS)} or 'all'"
        )
    failures: dict[str, str] = {}
    for name in selected:
        checker = _OPCHECK_TESTS[name]
        try:
            checker(op_def, args, kwargs)
        except Exception as exc:  # noqa: BLE001 - reported, not raised
            failures[name] = f"{type(exc).__name__}: {exc}"
    if failures and raise_exception:
        rendered = "\n".join(f"  {k}: {v}" for k, v in failures.items())
        raise RuntimeError(
            f"opcheck({op_def.name}) failed {len(failures)} check(s):\n"
            f"{rendered}"
        )
    return failures

[docs]
class Library:
    """

    one DEF library per process/namespace), ``"IMPL"`` adds kernels, and
    ``"FRAGMENT"`` extends an existing namespace from multiple locations.
    """

    def __init__(self, namespace: str, kind: str = "DEF") -> None:
        if not isinstance(namespace, str) or not namespace.isidentifier():
            raise ValueError(
                f"library namespace must be an identifier, got {namespace!r}"
            )
        if kind not in {"DEF", "IMPL", "FRAGMENT"}:
            raise ValueError(f"unknown Library kind {kind!r}")
        self.ns = namespace
        self.kind = kind
        self._op_names: list[str] = []
        with _LOCK:
            if kind == "DEF":
                if namespace in _DEFINED_LIBRARY_NAMESPACES:
                    raise RuntimeError(
                        f"only a single DEF Library may exist for namespace "
                        f"{namespace!r}"
                    )
                _DEFINED_LIBRARY_NAMESPACES.add(namespace)

    def _define(self, opname: str) -> CustomOpDef:
        full_name = f"{self.ns}::{opname}"
        with _LOCK:
            existing = _OP_REGISTRY.get(full_name)
        if existing is not None:
            if self.kind != "FRAGMENT" and existing.namespace == self.ns:
                raise RuntimeError(
                    f"operator {full_name!r} already defined in namespace "
                    f"{self.ns!r}"
                )
            return existing
        return _register_op_def(CustomOpDef(full_name))

[docs]
    def define(
        self, schema: str, *, alias_analysis: str = "", tags: Any = ()
    ) -> None:
        """Define an operator from a schema like ``"ns::add(Tensor, Tensor)"``.

        Only the qualified name is meaningful (TensorPlay models no schema
        schema strings can be pasted verbatim.  ``alias_analysis`` and
        ``tags`` are accepted and ignored for call-site compatibility.
        """

        del alias_analysis, tags
        if not isinstance(schema, str) or "::" not in schema:
            raise ValueError(
                f'schema must look like "ns::op(...)", got {schema!r}'
            )
        signature = schema.split("(", 1)[0].strip()
        namespace, opname = _validate_name(signature)
        if namespace != self.ns:
            raise ValueError(
                f"schema namespace {namespace!r} does not match library "
                f"namespace {self.ns!r}"
            )
        op_def = self._define(opname)
        if op_def.name not in self._op_names:
            self._op_names.append(op_def.name)

[docs]
    def impl(
        self,
        op_name: str,
        fn: Callable[..., Any] | None = None,
        *,
        device_type: str = "CompositeExplicitAutograd",
        dispatch_key: str = "",
        allow_override: bool = True,
    ) -> Callable[..., Any]:
        """

        ``device_type`` accepts composite spellings (``Composite…`` → the
        device-agnostic slot) or concrete devices (``"CPU"``/``"CUDA"``);
        (non-empty wins).  May be used directly or as a decorator.
        """

        if dispatch_key:
            device_type = dispatch_key
        if not isinstance(device_type, str):
            raise TypeError(
                f"device_type must be a str, got {type(device_type)!r}"
            )
        del allow_override
        if device_type in _COMPOSITE_KEYS:
            key: Any = None
        else:
            key = device_type.lower()
        target = f"{self.ns}::{op_name}" if "::" not in op_name else op_name
        op_def = _resolve_op(target)
        if fn is not None and not callable(fn):
            raise TypeError(f"kernel must be callable, got {type(fn)!r}")

        def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
            op_def.register_kernel(key)(func)
            return func

        return wrapper(fn) if fn is not None else wrapper

[docs]
    def fallback(self, kind: str) -> None:
        """

        TensorPlay's dispatcher has no per-key fallthrough table, so this
        of silently mis-dispatching.
        """

        raise NotImplementedError(
            "Library.fallback is not supported: TensorPlay's dispatcher has "
            "no per-dispatch-key fallthrough table (kind="
            f"{kind!r})"
        )

    def __repr__(self) -> str:
        return f"<Library ns={self.ns!r} kind={self.kind}>"
```