# tensorplay.library API

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

## Functions 14

[#](#api-tensorplay.library.custom_op)

### custom_op

function[Full reference ↗](/docs/generated/tensorplay.library.custom_op.html)

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

Example:

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

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

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

Parameters:

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

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

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

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

- schema – Optional schema string kept for introspection and [opcheck()](/docs/generated/tensorplay.library.opcheck.html#tensorplay.library.opcheck) (TensorPlay models no schema grammar).

Returns:

A decorator producing a callable CustomOpDef.

[#](#api-tensorplay.library.define)

### define

function[Full reference ↗](/docs/generated/tensorplay.library.define.html)

```python
tensorplay.library.define(qualname: str, schema: str | None = None, *, lib: Any = None, tags: Any = ()) → None
```

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

[#](#api-tensorplay.library.get_kernel)

### get_kernel

function[Full reference ↗](/docs/generated/tensorplay.library.get_kernel.html)

```python
tensorplay.library.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
CustomOpDef.set_kernel_enabled() visibility).

[#](#api-tensorplay.library.impl_abstract)

### impl_abstract

function[Full reference ↗](/docs/generated/tensorplay.library.impl_abstract.html)

```python
tensorplay.library.impl_abstract(qualname: str, func: Callable[[...], Any] | None = None, /, *, lib: Any = None) → Callable[[Callable[[...], Any]], Callable[[...], Any]]
```

[#](#api-tensorplay.library.impl)

### impl

function[Full reference ↗](/docs/generated/tensorplay.library.impl.html)

```python
tensorplay.library.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).

[#](#api-tensorplay.library.infer_schema)

### infer_schema

function[Full reference ↗](/docs/generated/tensorplay.library.infer_schema.html)

```python
tensorplay.library.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.

[#](#api-tensorplay.library.opcheck)

### opcheck

function[Full reference ↗](/docs/generated/tensorplay.library.opcheck.html)

```python
tensorplay.library.opcheck(op: str | CustomOpDef, args: tuple[Any, ...], kwargs: dict[str, Any] | None = None, *, test_utils: str | Sequence[str] = ('test_schema', 'test_autograd_registration', 'test_faketensor', 'test_aot_dispatch_dynamic'), raise_exception: bool = True, atol: float | None = None, rtol: float | None = None) &#x2192; 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.

[#](#api-tensorplay.library.register_autocast)

### register_autocast

function[Full reference ↗](/docs/generated/tensorplay.library.register_autocast.html)

```python
tensorplay.library.register_autocast(op: str | CustomOpDef, device_type: str, cast_inputs: Any, /, *, lib: Any = None) &#x2192; None
```

[#](#api-tensorplay.library.register_autograd)

### register_autograd

function[Full reference ↗](/docs/generated/tensorplay.library.register_autograd.html)

```python
tensorplay.library.register_autograd(op: str | CustomOpDef, backward: Callable[[...], Any], /, *, setup_context: Callable[[...], Any] | None = None, lib: Any = None) &#x2192; None
```

[#](#api-tensorplay.library.register_fake)

### register_fake

function[Full reference ↗](/docs/generated/tensorplay.library.register_fake.html)

```python
tensorplay.library.register_fake(op: str | CustomOpDef, func: Callable[[...], Any] | None = None, /, *, lib: Any = None, allow_override: bool = True) &#x2192; Callable[[Callable[[...], Any]], Callable[[...], Any]]
```

[#](#api-tensorplay.library.register_kernel)

### register_kernel

function[Full reference ↗](/docs/generated/tensorplay.library.register_kernel.html)

```python
tensorplay.library.register_kernel(op: str | CustomOpDef, device_types: Any = None, func: Callable[[...], Any] | None = None, /, *, lib: Any = None) &#x2192; Callable[[Callable[[...], Any]], Callable[[...], Any]]
```

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

[#](#api-tensorplay.library.register_vmap)

### register_vmap

function[Full reference ↗](/docs/generated/tensorplay.library.register_vmap.html)

```python
tensorplay.library.register_vmap(op: str | CustomOpDef, func: Callable[[...], Any] | None = None, /, *, lib: Any = None) &#x2192; Callable[[Callable[[...], Any]], Callable[[...], Any]]
```

[#](#api-tensorplay.library.triton_op)

### triton_op

function[Full reference ↗](/docs/generated/tensorplay.library.triton_op.html)

```python
tensorplay.library.triton_op(name: str, fn: Callable[[...], Any] | None = None, /, *, mutates_args: Sequence[str] = (), device_types: Any = None, schema: str | None = None) &#x2192; Callable[[Callable[[...], Any]], CustomOpDef]
```

The registered kernel(s) must launch their Triton kernels through
[wrap_triton()](/docs/generated/tensorplay.library.wrap_triton.html#tensorplay.library.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

[#](#api-tensorplay.library.wrap_triton)

### wrap_triton

function[Full reference ↗](/docs/generated/tensorplay.library.wrap_triton.html)

```python
tensorplay.library.wrap_triton(kernel: Any) &#x2192; TritonKernelWrapper
```

Mark a Triton kernel as launchable from within a triton_op.

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

## Classes 1

[#](#api-tensorplay.library.Library)

### Library

class[Full reference ↗](/docs/generated/tensorplay.library.Library.html)

```python
class tensorplay.library.Library(namespace: str, kind: str = 'DEF')
```

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

```python
define(schema: str, *, alias_analysis: str = '', tags: Any = ()) &#x2192; 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.

```python
fallback(kind: str) &#x2192; None
```

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

```python
impl(op_name: str, fn: Callable[[...], Any] | None = None, *, device_type: str = 'CompositeExplicitAutograd', dispatch_key: str = '', allow_override: bool = True) &#x2192; 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.