Skip to content

← Dtypes and Numerics step 3 of 6

Easy Primitives

Group integer indices with floor division

Jump to the task

Grouping integer indices requires a choice of rounding direction. For a positive group size, floor division assigns every index to the interval that contains it, including indices below zero.

import torch

idx = torch.tensor([-7, -1, 0, 1, 7], dtype=torch.int64)
print((idx // 2).tolist())  # [-4, -1, 0, 0, 3]

Group -1 contains indices -2 and -1; group 0 contains 0 and 1. The group number is a mathematical label here, not a request to index a Python list from its end.

Floor division versus truncation

Flooring rounds toward negative infinity. Truncation rounds toward zero. They differ for negative quotients that are not already integers:

print(torch.div(idx, 2, rounding_mode="floor").tolist())
# [-4, -1, 0, 0, 3]
print(torch.div(idx, 2, rounding_mode="trunc").tolist())
# [-3, 0, 0, 0, 3]

For integer operands, PyTorch's / performs true division and returns a floating-point result. Casting that result to an integer truncates toward zero, so it does not implement the required floor operation. Converting large integers through a floating dtype can also lose precision.

Use // or torch.div(..., rounding_mode="floor") directly on the integer inputs. With floating-point inputs, // can return a floating tensor, so do not infer its dtype from the operator alone.

Your task

def split_index(idx: torch.Tensor, group_size: int) -> torch.Tensor

idx is a one-dimensional int64 tensor and group_size is a positive integer. Return an int64 tensor of the same shape, where each output is floor(idx[i] / group_size). Handle negative indices as well as positive ones. The starter divides in floating point and then truncates.