Tensor Shape and Attributes
Problem Statement
Given a PyTorch tensor, extract and return its key attributes.
Background
Understanding tensor attributes is crucial for debugging and ensuring tensors are compatible for operations. Every tensor has a shape, data type, and device.
Your Task
Write a function get_tensor_info(tensor) that returns a dictionary containing the tensor's shape (as a list), dtype (as a string), and device (as a string).
Output Format
Return a dictionary with keys: "shape", "dtype", "device".
Example:
torch.rand(2, 3)
{"shape": [2, 3], "dtype": "torch.float32", "device": "cpu"}We access tensor.shape, tensor.dtype, and tensor.device attributes
Constraints:
- Input will always be a valid PyTorch tensor
- Tensor can be of any shape and data type
- Return dtype and device as strings
1. Background Knowledge
PyTorch tensors are the fundamental data structure for deep learning, analogous to NumPy arrays but with GPU acceleration and autograd support. Key attributes include:
- Shape: A tuple of integers representing dimensions (e.g., [3, 4] for a 3Ć4 matrix). Accessed via tensor.shape or tensor.size(). Shape mismatches cause runtime errors in operations like matrix multiplication.
- dtype: Data type of elements (e.g., torch.float32, torch.int64). Accessed via tensor.dtype. Critical for numerical stability and memory usage.
- device: Storage location (e.g., 'cpu', 'cuda:0'). Accessed via tensor.device. Ensures compatibility for operations across hardware.
These attributes enable debugging tensor operations in neural networks, where shape inference is undecidable in general but vital for correctness. Tensor diagrams from tensor network literature illustrate shapes visually:
Tensor shape [B, C, H, W] (batch, channels, height, width):
āāāāāāāāāāāāāāā
ā [B, C, H, W] ā ā Multi-dimensional array
āāāāāāāāāāāāāāā
2. Algorithm Approach
No complex algorithms needed; this is direct attribute extraction. The approach mirrors tensor inspection in PyTorch debugging tools and static analyzers:
- Query tensor properties using built-in accessors.
- Convert to required formats (tuple ā list for shape, dtype/device objects ā strings).
Pseudocode Flowchart (text-based, inspired by tensor network diagrams):
Input: tensor
ā
tensor.shape ā list(tensor.shape)
ā
tensor.dtype ā str(tensor.dtype)
ā
tensor.device ā str(tensor.device)
ā
Return: {"shape":..., "dtype":..., "device":...}
Time: O(1) per attribute (constant-time access).
3. Step-by-Step Strategy
- Import PyTorch: import torch.
- Define function get_tensor_info(tensor) taking a torch.Tensor.
- Extract shape: shape = list(tensor.shape) (convert tuple to list).
- Extract dtype: dtype = str(tensor.dtype) (e.g., 'torch.float32').
- Extract device: device = str(tensor.device) (e.g., 'cpu' or 'cuda:0').
- Return dictionary: return {"shape": shape, "dtype": dtype, "device": device}.
Complete Solution:
import torch
def get_tensor_info(tensor):
return {
"shape": list(tensor.shape),
"dtype": str(tensor.dtype),
"device": str(tensor.device)
}
Test:
t = torch.rand(3, 4)
print(get_tensor_info(t)) # {"shape": [3, 4], "dtype": "torch.float32", "device": "cpu"}
4. Common Pitfalls
- Shape as tuple: tensor.shape is a tuple; use list() for list output.
- dtype/device as objects: tensor.dtype is torch.dtype; str() converts to string like "torch.float32".
- Device on CPU-only systems: Returns "cpu"; no error if no GPU.
- Empty tensors: Shape [] works fine.
- Assuming input validity: Per constraints, no need for isinstance(tensor, torch.Tensor) checks.
- GPU tensors: str(tensor.device) handles 'cuda:0' correctly.
5. Time & Space Complexity
- Time: O(1) ā Attribute access is constant time, independent of tensor size.
- Space: O(d) where d is number of dimensions (storing shape list); negligible for practical tensors. Dictionary creation is O(1). No tensor copying occurs.