helion.language.register_reduction_dim

helion.language.register_reduction_dim(size)

Explicitly register a reduction dimension that should be used for reduction operations.

This is useful when you need to allocate a dimension for reduction that isn’t automatically inferred from a slice operation. The registered dimension can be used for allocations and operations that require knowing the reduction size upfront.

Parameters:

size (int) – An integer representing the reduction dimension size.

Returns:

A SymInt object representing the reduction dimension size.

Return type:

torch.SymInt