Shape & indexing

9 functions — 9 documented.

cat¶

Join tensors along an existing axis.

The output index is split into (outer, along, inner) around the joined axis: inner is the product of the extents after it, outer everything before, so an element's position along the axis is (i / inner) % extent. Positions below the first operand's extent come from it; the rest from the second with the offset subtracted.

Each source index is rebuilt with that source's own axis extent. Reusing the output's extent would read past the end of the shorter operand — a trap documented in both kernels, and the reason they compute it the same way.

That property is also what brought cat's push-constant block inside the 128-byte guarantee. The operands' extents were 32 bytes restating what the output's extents plus the joined axis already said, so the shader now reconstructs them and the block dropped from 144 to 112 bytes.

Parameters

tensors (Sequence[Tensor])

Tensors to join. All must agree on every axis except the joined one.

axis (int)

Axis to join along.

Returns

A new tensor whose extent along axis is the sum of the inputs'.

Example

copy>>> a = vkml.tensor(np.array([[1.0, 2.0]], dtype=np.float32))
>>> b = vkml.tensor(np.array([[3.0, 4.0]], dtype=np.float32))
>>> vkml.cat([a, b], 0).numpy()
array([[1., 2.],
       [3., 4.]], dtype=float32)

From the header

Joins tensors along axis. Every input must share the same rank, dtype, device and extents on every axis except axis.

The graph node is binary, so a list of N is folded left into N-1 nodes. That copies O(N^2) bytes for equal-sized inputs where O(N) is achievable with an n-ary node. Deferred deliberately: N is 2 for every use in scope (residual joins, UNet skips), and an n-ary node needs per-source extents in push constants for a case that does not yet exist.

In the CPU kernel

Joins two operands along one axis.

The linear index is split into (outer, along, inner) around the concatenated axis: inner is the product of the extents after it, outer everything before, so an element's position along the axis is (i / inner) % extent. Positions below the first operand's extent come from it, the rest from the second with the offset subtracted.

Each source index is rebuilt using the source's own axis extent, which is what makes this correct when the two differ -- reusing the output's extent would read past the end of the shorter operand.

Byte-wise rather than typed: concatenation moves elements without interpreting them, so one implementation serves every dtype.

src/backend/cpu/kernels_movement.cpp:275

Implementation

See also index_select, where

tril¶

Zero everything above the k-th diagonal.

Keeps the lower triangle. k=0 is the main diagonal, positive k keeps more above it, negative less.

Parameters

input (Tensor)

Rank 2 or higher; the last two axes are the matrix.

k (int = 0) optional

Diagonal offset.

Returns

A tensor of the same shape and dtype.

Example

copy>>> x = vkml.tensor(np.ones((3, 3), dtype=np.float32))
>>> vkml.tril(x, 0).numpy()
array([[1., 0., 0.],
       [1., 1., 0.],
       [1., 1., 1.]], dtype=float32)

From the header

Mirror of triu, keeping the lower triangle.

Implementation

See also triu, masked_fill

triu¶

Zero everything below the k-th diagonal.

Keeps the upper triangle. Combined with masked_fill this is how causal attention masks are built — nn.MultiheadAttention's is_causal path uses exactly that pair.

Parameters

input (Tensor)

Rank 2 or higher; the last two axes are the matrix.

k (int = 0) optional

Diagonal offset.

Returns

A tensor of the same shape and dtype.

Example

copy>>> x = vkml.tensor(np.ones((3, 3), dtype=np.float32))
>>> vkml.triu(x, 1).numpy()
array([[0., 1., 1.],
       [0., 0., 1.],
       [0., 0., 0.]], dtype=float32)

From the header

Zeroes everything below the diagonal-th diagonal of the last two axes, keeping the upper triangle. diagonal = 0 keeps the main diagonal; positive values move the boundary toward the upper-right. Rank must be at least 2; leading axes are batched over.

Implementation

See also tril, masked_fill, where

masked_fill¶

Replace elements where a bool mask is true with a constant.

The fill value travels as a push constant, so no tensor is allocated for it — which is the difference from where, where both branches are tensors that already exist.

Parameters

input (Tensor)

Any shape, float dtype.

mask (Tensor)

A bool tensor, broadcastable against input.

value (float)

What to write where the mask is true.

Returns

A tensor of the same shape and dtype as input.

Example

copy>>> x = vkml.tensor(np.array([1.0, 2.0, 3.0], dtype=np.float32))
>>> m = vkml.greater(x, vkml.tensor(np.array([1.5, 1.5, 1.5], dtype=np.float32)))
>>> vkml.masked_fill(x, m, 0.0).numpy()
array([1., 0., 0.], dtype=float32)

From the header

Replaces every element where mask is true with value, following torch.masked_fill. mask must be Bool and broadcastable to a's shape.

Composed from where rather than given its own kernel: the replacement is a rank-0 tensor broadcast with stride 0, so the "fused" version would save one graph node and four bytes. See the commit that added it.

Implementation

See also where, triu, clamp

index_select¶

Gather slices along one axis by index.

The index tensor must be rank 1 and int64. The output takes the index's length along the selected axis and keeps every other axis unchanged.

Parameters

input (Tensor)

Source tensor.

axis (int)

Axis to index along.

index (Tensor)

Rank-1 int64 indices.

Returns

A new tensor with index.numel() entries along axis.

Example

copy>>> x = vkml.tensor(np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32))
>>> idx = vkml.tensor(np.array([2, 0], dtype=np.int64))
>>> vkml.index_select(x, 0, idx).numpy()
array([[5., 6.],
       [1., 2.]], dtype=float32)

From the header

Gathers along axis at the positions named by index, following torch.index_select. index must be I64 and rank 1; the result takes its extent on axis from the index length, and every other extent from a.

This is Embedding's forward pass.

In the CPU kernel

Gathers rows named by an index vector. Embedding's forward pass.

src/backend/cpu/kernels_movement.cpp:339

Implementation

See also scatter_add, cat

scatter_add¶

Accumulate rows of a source into a zero tensor at given indices.

The inverse of index_select, and the operation its gradient needs. Several source rows may target the same destination, so the contributions must be added, not written — which is what makes this irreducible to the element-wise and reduction operators.

Both backends walk the source in ascending linear order, so for any destination the contributions arrive in ascending index order. That fixed order is what makes the result bit-reproducible, and the two backends agree exactly rather than merely within a tolerance.

Parameters

src (Tensor)

Source rows.

axis (int)

Axis to scatter along.

index (Tensor)

Rank-1 int64 indices, one per source row.

dim_size (int)

Extent of the output along axis.

Returns

A new tensor with dim_size entries along axis.

Example

copy>>> src = vkml.tensor(np.array([[1.0], [2.0], [3.0]], dtype=np.float32))
>>> idx = vkml.tensor(np.array([0, 0, 1], dtype=np.int64))
>>> vkml.scatter_add(src, 0, idx, 2).numpy()
array([[3.],
       [3.]], dtype=float32)

From the header

Adjoint of index_select: accumulates each slice of src into the row of a zero-filled result named by index, following torch.index_add. dim_size is the extent of the result on axis.

One of the four operations docs/ARCHITECTURE.md records as genuinely needing its own kernel rather than composing from forward ops -- repeated indices mean several source slices land on one destination, which no elementwise or reduction op expresses.

DETERMINISM. The target GPU has no global float atomicAdd (shaderBufferFloat32AtomicAdd = false, measured), so the usual atomic scatter is unavailable -- and unwanted, since its order varies run to run. Both backends instead fold contributions in ascending index order, which makes the result bit-reproducible and identical between them.

In the CPU kernel

Adjoint of index_select: accumulate each source slice into the row it came from. Repeated indices mean several sources land on one destination, which is what makes this irreducible to the elementwise and reduction ops.

Walks the SOURCE in ascending linear order, so for any destination the contributions arrive in ascending index order. That fixed order is what makes the result bit-reproducible, and it is the same order the Vulkan kernel uses -- deliberately, so the two agree exactly rather than merely within a tolerance. Fixing it costs nothing here and is the only way to have it at all on the GPU, which has no global float atomicAdd.

src/backend/cpu/kernels_movement.cpp:381

Implementation

See also index_select

im2col¶

Expand sliding windows into columns, the lowering behind conv2d.

(N, C, H, W) becomes (N, C·kh·kw, L), where L is the number of window positions. Every element of every window is written out, so a convolution becomes a single matrix multiply.

Parameters

input (Tensor)

(N, C, H, W).

kernel (Sequence[int])

Window size in (H, W).

stride (Sequence[int] = [1, 1]) optional

Step.

padding (Sequence[int] = [0, 0]) optional

Zero padding on both sides.

dilation (Sequence[int] = [1, 1]) optional

Spacing between window elements.

Returns

(N, C·kh·kw, L).

Example

copy>>> x = vkml.tensor(np.arange(16, dtype=np.float32).reshape(1, 1, 4, 4))
>>> vkml.im2col(x, [2, 2], [2, 2], [0, 0]).shape
(1, 4, 4)

From the header

Extracts sliding local blocks, following torch.nn.functional.unfold.

(N, C, H, W) becomes (N, C * kernel_h * kernel_w, L), where L is the number of window positions. Positions falling outside the padded image contribute zero. This is the first half of convolution-as-GEMM: the second half is an ordinary matmul against the flattened weights.

Implementation

See also col2im, conv2d, matmul

col2im¶

Fold columns back into an image, accumulating overlaps.

The adjoint of im2col, and the operation conv2d's gradient with respect to its input needs. Where windows overlap, the contributions are summed — which is what makes it the adjoint rather than merely the inverse.

Parameters

cols (Tensor)

(N, C·kh·kw, L).

image (Sequence[int])

Spatial size (H, W) to fold back into.

kernel (Sequence[int])

Window size in (H, W).

stride (Sequence[int] = [1, 1]) optional

Step.

padding (Sequence[int] = [0, 0]) optional

Padding used by the forward im2col.

dilation (Sequence[int] = [1, 1]) optional

Spacing between window elements.

Returns

(N, C, H, W).

From the header

Adjoint of im2col, following torch.nn.functional.fold: sums every window contribution back into the image position it came from.

Overlapping windows mean one image position receives several contributions, which is why this cannot be expressed as a gather and is one of the few operations needing its own kernel. image gives the spatial extent to reconstruct, which the column tensor does not determine.

Implementation

See also im2col, conv2d, scatter_add

detach¶

A tensor sharing the same values but outside the autograd graph.

Returns a tensor with requires_grad false and no recorded history, so a gradient cannot flow back through it. Used to freeze part of a model, or to stop a running statistic from being differentiated.

Parameters

input (Tensor)

The tensor to detach.

Returns

A tensor with the same values and no graph history.

Example

copy>>> x = vkml.tensor(np.array([1.0], dtype=np.float32), requires_grad=True)
>>> vkml.detach(x).requires_grad
False

From the header

A tensor sharing the same data but with no gradient history.

Implementation

See also backward, realize