Choosing the kind of custom op

When defining a custom op, first decide what aliasing and mutation behavior the operator has. This choice determines the mutates_args, tags, return schema, and fake-kernel behavior.

Does the operator mutate any Tensor input?
  |
  +-- no --> Functional operator
  |         Use mutates_args=().
  |         Return new Tensor objects that do not alias inputs or each other.
  |
  +-- yes --> Does it mutate the first positional Tensor argument and return it?
        |
        +-- yes --> In-place operator
        |          Use tags=tensorplay.Tag.inplace.
        |          Use mutates_args with exactly the first argument name.
        |          Return that same first argument.
        |
        +-- no --> Does it mutate keyword-only Tensor out= arguments and return them?
              |
              +-- yes --> out= operator
              |          Use tags=tensorplay.Tag.out.
              |          Do not read from the out= tensors.
              |          Put mutable output tensors after * as keyword-only args.
              |          Return all out= tensors in declaration order.
              |
              +-- no --> Mutable operator
                         Use mutates_args for the mutated arguments.
                         Do not return mutated inputs or aliases of inputs.

Functional operators are the simplest and should be preferred when possible. They should not mutate inputs, and their outputs should be fresh values that do not alias any input or other output. Use tags=tensorplay.Tag.inplace only for the conventional in-place shape: the first positional argument is a Tensor, it is the only mutated argument, and the operator returns that same Tensor. For example, an inferred in-place custom op has a schema shaped like:

(Tensor(a!) x, ...) -> Tensor(a!)

Use tags=tensorplay.Tag.out only for conventional out= variants: all mutable output tensors are keyword-only arguments, and the operator returns those outputs in declaration order. The implementation must write to out= tensors, but must not read from them. For example:

(Tensor x, *, Tensor(a!) out) -> Tensor(a!)
(Tensor x, *, Tensor(a!) out0, Tensor(b!) out1) -> (Tensor(a!), Tensor(b!))

In-place and out= custom ops get an autogenerated fake kernel from their tag and schema. Registering a fake kernel for these ops is usually unnecessary, though it is still allowed if you need to override the autogenerated behavior. out= custom ops follow the same autograd rule as built-in out= operators: they do not support automatic differentiation when any argument requires grad. Arbitrary aliasing, such as returning an alias of an input that is not one of the in-place or out= patterns above, is not supported by custom_op. TensorPlay compiler transforms such as functionalization need to reason precisely about mutations and aliases, so custom op aliasing is intentionally limited to the conventional forms described here. If your operator mutates inputs but is not an in-place or out= operator, model it as a mutable operator and return no aliases of inputs. The same operation can have different contracts depending on how it handles its outputs:

from tensorplay import Tensor


@tensorplay.library.custom_op(
    "mylib::add_inplace",
    mutates_args={"x"},
    tags=tensorplay.Tag.inplace,
)
def add_inplace(x: Tensor, y: Tensor) -> Tensor:
    x.add_(y)
    return x


@tensorplay.library.custom_op(
    "mylib::add_out",
    mutates_args={"out"},
    tags=tensorplay.Tag.out,
)
def add_out(x: Tensor, y: Tensor, *, out: Tensor) -> Tensor:
    out.copy_(x + y)
    return out


@tensorplay.library.custom_op("mylib::add_mutate", mutates_args={"x"})
def add_mutate(x: Tensor, y: Tensor) -> None:
    x.add_(y)