Reduction

7 functions — 7 documented.

`sum`¶

sum(input: Tensor, dim: object | None = None, keepdim: bool = False) → Tensor

CPUVulkan

Sum every element of a tensor.

Reduces to a 0-d tensor.

The fold is pairwise, not sequential: pairwise_sum in src/backend/cpu/reduce.h recurses until a run is at most kPairwiseBlock = 32 elements and only then adds in order. The error grows as O(log n) in the element count rather than O(n), which is what keeps a large reduction usable in float32 at all.

On the GPU the same shape is achieved differently: each invocation folds its own strided slice, then the workgroup combines through a shared-memory tree. Both trees are determined by the tensor's shape, never by which invocation finished first.

Parameters

input (Tensor): Any shape, float dtype.

Returns

A 0-d tensor holding the total.

ⓘ Note

Bit-identical across runs on the same device, and that is a consequence of the fixed tree rather than a coincidence — an atomic-accumulation reduction would give a different answer each run, which is why vkML does not use one.

Example

>>> x = vkml.tensor(np.ones((1000,), dtype=np.float32))
>>> float(vkml.sum(x).item())
1000.0

Implementation

Declared in	`include/vkml/api/ops.h:377`
Graph node	`OpKind::Sum`
CPU kernel	`src/backend/cpu/kernels_reduce.cpp:94`
Vulkan shader	`shaders/reduce.comp` (286 lines · 4 specialisation constants)
Gradient rule	`autograd.cpp:462`
Decisions	ADR 0010
Benchmarked	`sum` · `sum 1024x1024 [`
History	commits touching the CPU kernel
Tests (≥70)	`test_autograd_vs_torch.py` `test_backend_parity.py` `test_device_limits.py` `test_f16.py` `test_invariants.py` `test_layout_and_scale.py` `test_nan_semantics.py` `test_nn_vs_torch.py` `test_ops_vs_torch.py` `test_vulkan_kernels.py`

See also mean, prod, amax, amin

`mean`¶

mean(input: Tensor, dim: object | None = None, keepdim: bool = False) → Tensor

CPUVulkan

Arithmetic mean of every element.

Computed as pairwise_sum(x) / n — the division happens once, at the end, rather than accumulating x/n per element. Dividing first would lose the low bits of every term before adding them.

Parameters

input (Tensor): Any shape, float dtype.

Returns

A 0-d tensor holding the mean.

Example

>>> vkml.mean(vkml.tensor(np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32))).item()
2.5

Implementation

Declared in	`include/vkml/api/ops.h:378`
Graph node	`OpKind::Mean`
CPU kernel	`src/backend/cpu/kernels_reduce.cpp:98`
Vulkan shader	`shaders/reduce.comp` (286 lines · 4 specialisation constants)
Gradient rule	`autograd.cpp:482`
Decisions	ADR 0003 ADR 0010
History	commits touching the CPU kernel
Tests (≥6)	`test_f16.py` `test_layout_and_scale.py` `test_nn_vs_torch.py` `test_ops_vs_torch.py` `test_vulkan_kernels.py`

See also sum, amax

`prod`¶

prod(input: Tensor, dim: object | None = None, keepdim: bool = False) → Tensor

CPUVulkanCPU-only by decision

The product of every element.

Folded sequentially, in index order, and deliberately so. Every other reduction here folds pairwise because that improves a sum's error bound. A product gains nothing from it — relative errors compose multiplicatively whatever the order — and reordering costs something real: it changes when the fold overflows.

Multiplying 1e20 and 1e-20 alternately stays at 1.0 in index order and reaches inf if the large values are grouped together first.

Parameters

input (Tensor): Any shape, float dtype.

Returns

A 0-d tensor holding the product.

ⓘ Note

It also has no gradient rule, so backward through a prod raises.

⚠ Warning

CPU only, and the reason is numerical rather than an omission. A GPU reduction is a tree, and a tree reassociates the fold — which is exactly what changes the overflow point. Rather than ship a kernel that disagrees with the oracle on inputs like the one above, prod raises NotImplementedError on a Vulkan tensor. The rationale is recorded above k_prod in kernels_reduce.cpp.

Example

>>> vkml.prod(vkml.tensor(np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32))).item()
24.0

In the CPU kernel

Sequential, in index order, and deliberately so.

Every other reduction here folds pairwise, which improves a SUM's error bound. A product gains nothing from that -- relative errors compose multiplicatively whatever the order -- and reordering costs something real: it changes when the fold overflows. Multiplying 1e20 and 1e-20 alternately stays at 1.0 in index order and reaches inf if the large values are grouped.

That is why prod has no Vulkan kernel; see VulkanBackend::supports.

src/backend/cpu/kernels_reduce.cpp:118

Implementation

Declared in	`include/vkml/api/ops.h:379`
Graph node	`OpKind::Prod`
CPU kernel	`src/backend/cpu/kernels_reduce.cpp:118`
Vulkan shader	not implemented — raises NotImplementedError
Gradient rule	none — backward through it raises
Decisions	ADR 0008
History	commits touching the CPU kernel
Tests (≥5)	`test_backend_parity.py` `test_device_report.py` `test_layout_and_scale.py` `test_ops_vs_torch.py`

See also sum, mean

`amax`¶

amax(input: Tensor, dim: object | None = None, keepdim: bool = False) → Tensor

CPUVulkan

The largest element of a tensor.

Parameters

input (Tensor): Any shape, float dtype.

Returns

A 0-d tensor holding the maximum.

⚠ Warning

NaN propagates, matching torch.amax. The check is written twice in the Vulkan kernel — once in the per-invocation fold and once in the shared-memory tree — because a reduction built only from > comparisons drops NaN at whichever stage it is missing, and a NaN silently disappearing from a reduction hides a diverged model.

Example

>>> x = vkml.tensor(np.array([1.0, float('nan'), 3.0], dtype=np.float32))
>>> vkml.amax(x).item()
nan

Implementation

CPU kernel	composed from other operators
Vulkan shader	composed, or dispatched through a shared kernel
Gradient rule	none — backward through it raises
Tests (≥4)	`test_device_limits.py` `test_f16.py` `test_layout_and_scale.py` `test_vulkan_kernels.py`

See also amin, argmax, maximum, sum

`amin`¶

amin(input: Tensor, dim: object | None = None, keepdim: bool = False) → Tensor

CPUVulkan

The smallest element of a tensor.

Parameters

input (Tensor): Any shape, float dtype.

Returns

A 0-d tensor holding the minimum.

⚠ Warning

NaN propagates, checked in both fold stages, exactly as for amax.

Example

>>> vkml.amin(vkml.tensor(np.array([3.0, 1.0, 2.0], dtype=np.float32))).item()
1.0

Implementation

CPU kernel	composed from other operators
Vulkan shader	composed, or dispatched through a shared kernel
Gradient rule	none — backward through it raises
Tests (≥2)	`test_layout_and_scale.py` `test_vulkan_kernels.py`

See also amax, argmin, minimum

`argmax`¶

argmax(input: Tensor, dim: int, keepdim: bool = False) → Tensor

CPUVulkan

The index of the largest element along an axis.

Ties keep the first maximum. The comparison is a strict >, which the CPU kernel notes is what torch.argmax documents — >= would silently return the last one instead.

Parameters

input (Tensor): Any shape, float dtype.
dim (int): Axis to search along.
keepdim (bool = False) optional: Keep the reduced axis with extent 1.

Returns

An int64 tensor of indices.

Example

>>> x = vkml.tensor(np.array([[1.0, 5.0, 5.0]], dtype=np.float32))
>>> vkml.argmax(x, 1).numpy()
array([1], dtype=int64)

Implementation

Declared in	`include/vkml/api/ops.h:382`
CPU kernel	`src/backend/cpu/kernels_reduce.cpp:160`
Vulkan shader	`shaders/reduce.comp` (286 lines · 4 specialisation constants)
Gradient rule	none — backward through it raises
Decisions	ADR 0010
Benchmarked	`argmax`
History	commits touching the CPU kernel
Tests (≥3)	`test_f16.py` `test_layout_and_scale.py` `test_vulkan_kernels.py`

See also amax, argmin

`argmin`¶

argmin(input: Tensor, dim: int, keepdim: bool = False) → Tensor

CPUVulkan

The index of the smallest element along an axis.

Ties keep the first minimum, mirroring argmax.

Parameters

input (Tensor): Any shape, float dtype.
dim (int): Axis to search along.
keepdim (bool = False) optional: Keep the reduced axis with extent 1.

Returns

An int64 tensor of indices.

Example

>>> x = vkml.tensor(np.array([[3.0, 1.0, 1.0]], dtype=np.float32))
>>> vkml.argmin(x, 1).numpy()
array([1], dtype=int64)

Implementation

Declared in	`include/vkml/api/ops.h:383`
CPU kernel	`src/backend/cpu/kernels_reduce.cpp:178`
Vulkan shader	`shaders/reduce.comp` (286 lines · 4 specialisation constants)
Gradient rule	none — backward through it raises
Decisions	ADR 0010
History	commits touching the CPU kernel
Tests (≥1)	`test_layout_and_scale.py`

See also amin, argmax

Reduction

sum¶

Parameters

Returns

Example

Implementation

mean¶

Parameters

Returns

Example

Implementation

prod¶

Parameters

Returns

Example

In the CPU kernel

Implementation

amax¶

Parameters

Returns

Example

Implementation

amin¶

Parameters

Returns

Example

Implementation

argmax¶

Parameters

Returns

Example

Implementation

argmin¶

Parameters

Returns

Example

Implementation

`sum`¶

`mean`¶

`prod`¶

`amax`¶

`amin`¶

`argmax`¶

`argmin`¶