AxisReduction

The outer × axisSize × inner view of a tensor reduced along one axis, plus the result shape. Every backend's sumAxis / maxAxis / argMaxAxis / maxAxisBackward builds this same decomposition and feeds outer, axisSize, inner to its reduce_axis_* kernel; computing it here keeps the index math and the bounds check identical across CUDA, ROCm, and Vulkan.

View source ↗Full module reference on javadoc.io ↗

Record components

shape: Shape
axis: int
outer: int
axisSize: int
inner: int

Methods

static AxisReduction of(Shape shape, int axis)

Decomposes shape about axis, validating the axis is in range.

int outputSize()

Output element count: one per (outer, inner) pair.

Shape outputShape(boolean keepDims)

Result shape: axis kept as size 1 when keepDims, otherwise dropped.

void requireGradient(Shape gradientShape, int gradientCount)

Validates that an upstream-gradient buffer carries exactly one value per reduced slice, as maxAxisBackward requires.