While some frameworks allocate new tensors on the CPU by default, others, such as JAX and TensorFlow, utilize an accelerator-preferred allocation strategy. In these frameworks, newly created tensors are automatically stored on a GPU or TPU if one is detected on the system; they only fall back to using the CPU if no accelerated hardware is available.

Claude

By default, tensors are created in the main memory, and the system's CPU is used to perform calculations on them. However, practitioners can explicitly specify other devices, such as GPUs, for storage and calculation.

Default Tensor Device Allocation

Dive into Deep Learning

Deep learning frameworks provide programmatic mechanisms to determine the specific hardware device—such as main memory (CPU) or a specific GPU—where a tensor is currently stored. Depending on the framework, this information can be retrieved using attributes or methods on the tensor object, such as .device in PyTorch and TensorFlow, .device() in JAX, or .ctx in MXNet.

Querying Tensor Device Location

Accelerator-Preferred Default Tensor Allocation

Whenever a mathematical operation involves multiple tensors, such as calculating the sum of two tensors, it is required that all participating arguments live on the exact same hardware device. If an attempt is made to compute an operation across different devices—such as adding a CPU tensor to a GPU tensor—the framework will fail, because it lacks the necessary information to decide where the computation should be executed or where the resulting tensor should be stored in memory.

Learn Before

Related