Which half of the batch

a = torch.tensor([7, 8, 9])

a / 2      # tensor([3.5, 4.0, 4.5])   float32
a // 2     # tensor([3, 4, 4])         int64

/ is true division and always produces a float, even between two integer tensors. // floors and stays integral.

Where it matters

Anything that computes an index:

head = idx / head_dim        # float32. Indexing with it raises.
head = idx // head_dim       # int64. Correct.

A float index is not silently truncated; PyTorch refuses it. So this one fails loudly, which is the good case.

The one that does not fail loudly

// floors, it does not truncate toward zero:

 7 // 2  ->   3
-7 // 2  ->  -4       not -3

Python and PyTorch agree on this, and C, C++ and Rust do not: their integer division truncates toward zero, so -7 / 2 is -3. Porting a formula between the two without noticing gives an off-by-one that only appears for negative inputs.

If you want truncation, ask for it:

torch.div(a, 2, rounding_mode="trunc")     # toward zero
torch.div(a, 2, rounding_mode="floor")     # same as //

Naming the rounding mode is worth the extra characters wherever negatives are possible, because the two spellings look equally correct and differ only there.

Your task

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

Return which group each index falls into: index i belongs to group floor(i / group_size). The result must be an integer tensor, and must be right for negative indices.