Tensor Slicing and Indexing

mediumPROBasic OpsPyTorch OperationsPytorch

Problem Statement

Given a 2D tensor, extract specific rows, columns, and sub-matrices using PyTorch indexing.

Background

PyTorch tensors support NumPy-style indexing and slicing, allowing you to extract arbitrary parts of a tensor.

Your Task

Write a function extract_parts(tensor) that takes a 4×4 tensor and returns a dictionary with its first row, last column, and center 2×2 sub-matrix.

Output Format

Return a dictionary with keys: "first_row", "last_column", "center".

Each value must be a PyTorch tensor (first_row and last_column have 4 elements, center is 2×2). The sample below shows them as lists for readability.

Example

Input

torch.arange(16).reshape(4, 4)

Output

{"first_row": [0, 1, 2, 3], "last_column": [3, 7, 11, 15], "center": [[5, 6], [9, 10]]}

Reasoning

Input: a $4 \times 4$ integer tensor. In the sample its entries are $0$ to $15$ laid out row by row, so the entry in row $i$ , column $j$ (counting from $0$ ) is $t_{ij} = 4i + j$ .
first_row: fix $i = 0$ and let $j$ run from $0$ to $3$ : $t_{00}, t_{01}, t_{02}, t_{03} = 0, 1, 2, 3$ . A 1-D tensor of length 4.
last_column: fix $j = 3$ and let $i$ run from $0$ to $3$ : $4i + 3$ gives $3, 7, 11, 15$ . Each row down adds 4. Notice it comes back as a flat length-4 vector, not a $4 \times 1$ column.
center: drop the outer ring and keep rows $\{1, 2\}$ crossed with columns $\{1, 2\}$ : $t_{11} = 5$ , $t_{12} = 6$ , $t_{21} = 9$ , $t_{22} = 10$ , giving [[5, 6], [9, 10]]. This part stays 2-D.
Output: a dictionary mapping the three keys to tensors (shown as lists for display). An all-ones input gives all-ones parts with the same shapes.

Constraints

Input tensor is always 4x4
Values are integers
Return a dictionary with the three required keys

Learn this in

PyTorch FundamentalsFoundations

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 primerPREVIEW

1. Background Knowledge

PyTorch tensors are multi-dimensional arrays similar to NumPy arrays, supporting advanced indexing and slicing operations that enable efficient data extraction without copying data (views are returned when possible). Key concepts include:

Basic Indexing: tensor[i, j] accesses element at row i, column j.
Slicing: Uses Python slice notation start:stop:step.
Rows: tensor (first row), tensor[-1] (last row).
Columns: tensor[:, -1] (last column, : selects all rows).
Sub-matrices: tensor[row_slice, col_slice].
Tensor Shape: For a 4×4 tensor, shape = (4, 4). Center 2×2 is rows/cols [1:3, 1:3] (indices 1,2).
Views vs. Copies: Slicing creates views (zero-copy), preserving gradients for autograd.

These operations are $O(1)$ for view creation, essential for memory-efficient ML pipelines.

2. Algorithm Approach

This problem uses direct indexing/slicing, a fundamental tensor operation in PyTorch (NumPy-compatible). No complex algorithms needed:

First row: Single-index slice on dimension 0.
Last column: Full slice on dim 0, single index on dim 1.
Center sub-matrix: 2D slice on both dimensions.

For fixed-size (4×4) tensors, compute indices statically:

Last column index: tensor.size(1) - 1.
Center: (n//2-1): (n//2+1) where $n=4$ .

3. Step-by-Step Strategy

Inspect input: Confirm tensor.shape == (4, 4) (per constraints).
Extract first row: first_row = tensor.
Extract last column: last_col = tensor[:, -1] (or tensor[:, 3]).
Extract center 2×2: center = tensor[1:3, 1:3].
Package results: Return {"first_row": first_row, "last_column": last_col, "center": center}.

Complete Solution:

import torch

def extract_parts(tensor):
 """
 Extract specific parts from a 4x4 tensor.

 Args:
 tensor: 4x4 PyTorch tensor

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results