# 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() # Hot-path aliases resolved once at import (this module is imported last by # tensorplay/__init__, so every attribute below already exists). _is_grad_enabled = tensorplay.is_grad_enabled # 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() _validate_name(signature) 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 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() 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() 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: op_def = self class _CustomOpAutograd(tensorplay.autograd.Function): @staticmethod def forward(*args: Any, **kwargs: Any) -> Any: return op_def._run_kernel(args, kwargs) @staticmethod def setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None: if op_def._setup_context is not None: op_def._setup_context(ctx, inputs, 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 _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: key = _first_device_key(args) if self._autocast_rules: 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() } if _is_grad_enabled(): if _profiling_sessions: return self._profiled( self._autograd_cls.apply, args, kwargs) return self._autograd_cls.apply(*args, **kwargs) return self._run_profiled(args, kwargs, key) 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, 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 return self._eager_call(args, kwargs) def _first_device_key(values: tuple[Any, ...]) -> str | None: 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 ``((!))`` 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"" ```