Skip to content

← Devices and Data Movement step 2 of 4

Easy Primitives

Where does a new tensor go

Jump to the task

When a function creates a tensor to combine with its input, choose the new tensor's device and dtype deliberately. A factory call does not infer those properties from a tensor you plan to use later:

import torch

x = torch.ones(3, dtype=torch.float64)
noise = torch.randn_like(x)
print(noise.device == x.device)  # True
print(noise.dtype == x.dtype)    # True

By contrast, torch.randn(x.shape) uses the current defaults. The default device starts as CPU but can be changed. If x is on a different device, adding the two vectors can fail with a device mismatch.

Match the input when that is the intended type

For floating-point input x, these create normal noise on its device and with its dtype:

noise = torch.randn(x.shape, device=x.device, dtype=x.dtype)
noise = torch.randn_like(x)
noise = x.new_empty(x.shape).normal_()

The like factories reuse the input's shape by default. Tensor methods such as new_zeros let you specify a different shape while inheriting its device and dtype. Both families accept explicit overrides.

Matching dtype is not always the goal: an index tensor may need an integer dtype even when the data is floating-point, and a mask may need bool. Choose those types explicitly while keeping the required device.

Why the range's dtype matters here

With integer bounds, torch.arange(n) defaults to int64. Multiplying that range by a Python float produces a floating-point tensor using the default floating dtype. With the usual float32 default, adding it to float16 data promotes the result to float32. It is the floating multiplication and ensuing promotion, not an unconditional rule that any int64 range widens float16.

Creating this exercise's arithmetic range with dtype=x.dtype avoids that unintended change. The test vectors are short: low-precision floating ranges cannot represent every integer exactly once the indices become large.

Your task

def add_scaled_range(x: torch.Tensor, scale: float) -> torch.Tensor

For a one-dimensional floating-point input of length n, return x + scale * torch.arange(n). Create the range on x.device with x.dtype. The result must keep both properties.

For [1.0, 2.0, 3.0] and scale=2.0, the output values are [1.0, 4.0, 7.0]. The harness reports values rounded to three decimals, the dtype name and whether the output stayed on the input device.

The supplied cases expose the dtype issue on CPU. The explicit device choice also matters when this function is used with an accelerator input.