UX limitations

Operator coverage

vmap() executes func once over batched inputs, so every operator func calls needs a batching rule. An operator without one raises NotImplementedError naming the operator, for example:

vmap(lambda a, b: a.dot(b))(tp.randn(3, 4), tp.randn(3, 4))
# NotImplementedError: Kernel not found for op: dot on backend: VmapCPU

Verified to work under vmap: the arithmetic operators (+, -, *, /, **, unary -), the elementwise math functions (abs, exp, log, sqrt, sin, cos, tanh, sigmoid, relu), matmul, mm, bmm, sum, cumsum, clamp, where, stack, cat, index_select, narrow, select, tril, maximum, minimum, logsumexp, the shape operations (reshape, view, transpose, permute, movedim, squeeze, unsqueeze, expand, contiguous), basic and advanced indexing, the new_* factories plus tp.zeros/tp.ones, and randn.

Not yet covered — they raise the NotImplementedError above: dot, mean, max, min, argmax, softmax, log_softmax, norm, masked_fill, clone, sort, argsort, topk, einsum, dropout, randn_like, the norm layers (batch_norm, group_norm, layer_norm), nonzero, and item.

In-place operations

In-place arithmetic has no batching rule and raises cleanly:

vmap(lambda x, y: x.add_(y), in_dims=(0, 0))(tp.randn(3, 1), tp.randn(3, 1))
# NotImplementedError: Kernel not found for op: add_.Tensor on backend: VmapCPU

This holds whether the mutated tensor is batched or not — prefer the out-of-place form under vmap, or x + y directly.

Data-dependent operations

.item() raises RuntimeError under vmap — converting a batched value to a scalar is rejected because the result would depend on which sample is looked at — and nonzero raises NotImplementedError; the output of a data-dependent op varies per sample, so no single stacked result exists. Rewrite the code to avoid materializing per-sample shapes.

Data-dependent Python control flow

def relu(x):
    if x > 0:      # x is a mapped tensor here: do not do this
        return x
    return 0 * x

vmap(relu)(tp.randn(3))
# RuntimeError: item() is not supported on a vmap-batched tensor;
# data-dependent control flow cannot see per-batch values. ...

The out= keyword

The public operator signatures do not accept out= (passing one raises TypeError about the argument combination), so there is no out= form for vmap to transform.

Randomness

The intent of a random operation under vmap is ambiguous — should every sample see fresh values, or the same ones? — so vmap takes a randomness flag instead of guessing:

• "error" (the default): any random operation raises RuntimeError: random operations are not allowed under randomness=error.

• "different": elements of the batch draw different values.

• "same": every element of the batch sees the same value.

def add_noise(x):
    return x + tp.randn(())

x = tp.ones(3)
r = vmap(add_noise, randomness="same")(x)
assert (r == r[0]).all()          # one value, repeated

r = vmap(add_noise, randomness="different")(x)
assert (r != r[0]).any()          # fresh value per sample

An invalid value raises before the call: only "error", "different", and "same" are accepted. The flag governs the tensor factories such as tp.randn; dropout and randn_like do not have batching rules yet and raise regardless of the flag.