Access Dataset Sample

easyFREEDatasetsPyTorch DatasetsPytorch

Problem Statement

Given a PyTorch Dataset, extract a specific sample and return information about it.

Background

PyTorch Datasets support indexing — you can access individual samples, check the total number of samples, and iterate through the dataset.

Your Task

Write a function get_sample_info(dataset, index) that returns a dictionary with the total number of samples, the shape of the feature at the given index, and the label value at that index.

Output Format

Return a dictionary with keys: "num_samples" (int), "feature_shape" (list), "label_value" (int).

Example

Input

dataset with features [[1,2,3],[4,5,6]], labels [0,1], index=0

Output

{"num_samples": 2, "feature_shape": [3], "label_value": 0}

Reasoning

Inputs: dataset pairs a $2 \times 3$ feature matrix with a label vector of length 2, so it holds 2 samples. Sample $i$ is row $i$ of the features together with entry $i$ of the labels. index is 0, the first sample.
Sample count: there are 2 rows and 2 labels, so num_samples is 2. It counts samples, not the $2 \times 3 = 6$ individual feature values.
Feature shape: sample 0's feature is the row $[1, 2, 3]$ , a 1-D vector with 3 entries. Its shape is therefore [3]. It is not $[2, 3]$ , because that is the shape of the whole feature matrix, not of one sample.
Label: entry 0 of $[0, 1]$ is $0$ , so label_value is 0, reported as a plain integer rather than a one-element tensor.
Notice: num_samples doesn't depend on index. With index $= 1$ you would get the same count and shape, but the label would be $1$ .

Constraints

Index will be valid (within dataset bounds)
Return feature_shape as a list, not torch.Size
label_value should be a Python int, not tensor

More from Pytorch

Apply Softmax for Predictionseasy Count Model Parametersmedium Build with nn.Sequentialeasy Use nn.Flatten Layereasy Create a Simple Neural Network Classeasy

Concept primer

1. Background Knowledge

PyTorch Dataset and TensorDataset fundamentals:

PyTorch Dataset is an abstract class implementing len() and getitem() for iterable data access
TensorDataset combines multiple tensors into a dataset where dataset[i] returns a tuple (features[i], labels[i])
Indexing returns tuples of tensors - critical for unpacking correctly
len(dataset) returns the length of the shortest input tensor

Key tensor operations needed:

tensor.shape → tuple of dimensions
tensor.tolist() → Python native types
int(tensor.item()) → scalar tensor to Python int

2. Algorithm Approach

Direct indexing with tuple unpacking:

sample = dataset[index] # Returns (features_tensor, label_tensor)

No complex algorithms needed - pure O(1) direct access leveraging PyTorch's optimized tensor storage.

Three independent computations:

Dataset size: len(dataset) → O(1)
Feature shape extraction: sample.shape → O(1)
Label scalarization: int(sample.item()) → O(1)

3. Step-by-Step Strategy

def get_sample_info(dataset, index):
 # Step 1: Get total samples
 num_samples = len(dataset)

 # Step 2: Extract sample tuple
 sample = dataset[index] # (features, label)

 # Step 3: Extract feature shape as list
 feature_shape = list(sample.shape)

 # Step 4: Convert label tensor to Python int
 label_value = int(sample.item())

 # Step 5: Return exact dictionary format
 return {
 "num_samples": num_samples,
 "feature_shape": feature_shape,
 "label_value": label_value
 }

Execution flow:

len(dataset) → integer count
dataset[index] → tuple (features_tensor, label_tensor)
features_tensor.shape → torch.Size → list() conversion
label_tensor.item() → Python scalar → int() conversion

4. Common Pitfalls

Mistake	Why It Fails	Fix
return tensor.shape	Returns torch.Size, not list	list(tensor.shape)
return sample	Returns tensor, not Python int	int(sample.item())
features, labels = dataset[index]	Unpacks tuple correctly but shadows variable names	Use sample = dataset[index] then sample, sample
Missing .item() on scalar tensor	TypeError or tensor returned	Always use .item() for scalar extraction
sample.size()	Returns torch.Size	Use .shape property

Type conversion cheatsheet:

tensor.shape → torch.Size() → list(tensor.shape) →
scalar_tensor → tensor() → tensor.item() → 0.0
scalar_tensor → tensor() → int(tensor.item()) → 0

5. Time & Space Complexity

Time Complexity: $O(1)$

len(dataset): O(1) - stored metadata access
dataset[index]: O(1) - direct tensor slicing
Shape extraction: O(1) - metadata read
.item(): O(1) - scalar extraction

Space Complexity: $O(1)$

Only extracts one sample regardless of dataset size $N$
Temporary tuple and metadata only
Output dictionary: fixed 3 entries

Scalability:

Time: Independent of N (dataset size)
Space: Independent of N

Complete Working Solution

import torch
from torch.utils.data import TensorDataset

def get_sample_info(dataset, index):
 """Extract dataset info at specific index."""
 num_samples = len(dataset)
 sample = dataset[index]
 return {
 "num_samples": num_samples,
 "feature_shape": list(sample.shape),
 "label_value": int(sample.item())
 }

# Test
features = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
labels = torch.tensor([0, 1])
dataset = TensorDataset(features, labels)
print(get_sample_info(dataset, 0))
# {"num_samples": 2, "feature_shape":, "label_value": 0}

Key insight: This problem tests understanding of PyTorch Dataset indexing protocol and tensor-to-Python type conversion - fundamental for any PyTorch data pipeline.

solution.py

Python 3.13

Test results