helion.language.reduce

helion.language.reduce(combine_fn, input_tensor, dim=None, other=0, keep_dims=False)

Applies a reduction operation along a specified dimension or all dimensions.

This function is only needed for user-defined combine functions. Standard PyTorch reductions (such as sum, mean, amax, etc.) work directly in Helion without requiring this function.

Parameters:

• combine_fn (Union[Callable[[Tensor, Tensor], Tensor], Callable[..., tuple[Tensor, ...]]]) – A binary function that combines two elements element-wise. Must be associative and commutative for correct results. Can be tensor->tensor or tuple->tuple function.

• input_tensor (Tensor | tuple[Tensor, ...]) – Input tensor or tuple of tensors to reduce

• dim (int | None) – The dimension along which to reduce (None for all dimensions)

• other (float | tuple[float, ...]) – Value for masked/padded elements (default: 0) For tuple inputs, can be tuple of values with same length

• keep_dims (bool) – If True, reduced dimensions are retained with size 1

Returns:

Tensor(s) with reduced dimensions

Return type:

torch.Tensor or tuple[torch.Tensor, …]

See also

• associative_scan(): For prefix operations

Note

• combine_fn must be associative and commutative

• For standard reductions, use PyTorch functions directly (faster)

• Masked elements use the ‘other’ value during reduction