Matrix Multiplication and Element-wise Operations

mediumPROBasic OpsPyTorch OperationsPytorch

Problem Statement

Implement functions to perform both matrix multiplication and element-wise multiplication on tensors.

Background

PyTorch supports two types of multiplication that behave very differently:

Matrix multiplication: follows linear algebra rules (dot products of rows and columns)
Element-wise multiplication: multiplies corresponding elements directly

Understanding the difference is crucial for neural network operations.

Your Task

Write two functions:

matrix_multiply(a, b) — returns the matrix product of two tensors
elementwise_multiply(a, b) — returns the element-wise product of two tensors

Output Format

Each function should return a tensor.

Example

Input

a=[[1, 2], [3, 4]], b=[[5, 6], [7, 8]]

Output

{"matmul": [[19, 22], [43, 50]], "elementwise": [[5, 12], [21, 32]]}

Reasoning

Inputs: a and b are $2 \times 2$ integer matrices, $a = \begin{pmatrix}1 & 2\\3 & 4\end{pmatrix}$ and $b = \begin{pmatrix}5 & 6\\7 & 8\end{pmatrix}$ .
Matrix product: entry $(i, j)$ is row $i$ of $a$ dotted with column $j$ of $b$ , $c_{ij} = \sum_k a_{ik} b_{kj}$ . Row 1: $1 \cdot 5 + 2 \cdot 7 = 19$ , $1 \cdot 6 + 2 \cdot 8 = 22$ . Row 2: $3 \cdot 5 + 4 \cdot 7 = 43$ , $3 \cdot 6 + 4 \cdot 8 = 50$ . That gives [[19, 22], [43, 50]] under matmul.
Element-wise product: each entry is just $a_{ij} \cdot b_{ij}$ , with no summing: $5, 12, 21, 32$ , giving [[5, 12], [21, 32]] under elementwise.
Output: both results are $2 \times 2$ integer matrices, since integer products and sums stay integers.
Notice the shape rules differ: an $m \times n$ matrix times an $n \times p$ matrix gives $m \times p$ , while the element-wise product needs identical shapes. Swapping $a$ and $b$ changes the matrix product but not the element-wise one.

Constraints

For matrix multiplication, dimensions must be compatible
For element-wise multiplication, tensors must have same shape
Input tensors contain integers

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

Matrix multiplication and element-wise operations are foundational in ML/AI, powering linear layers, attention mechanisms, and convolutions in neural networks. PyTorch provides optimized implementations via dedicated operators.

Matrix Multiplication (MatMul): For matrices $A \in \mathbb{R}^{m \times n}$ and $B \in \mathbb{R}^{n \times p}$ , the result $C = A \cdot B$ has element $c_{ij} =$ \sum_{k=1} $^n a_{ik} b_{kj}$ . This follows linear algebra rules requiring compatible inner dimensions. PyTorch uses @ or torch.matmul(), leveraging cuBLAS/cuDNN for $O(mnp)$ complexity on GPUs.

Element-wise Multiplication (Hadamard Product): For same-shaped tensors $A, B \in \mathbb{R}^{m \times n}$ , $C = A \odot B$ has $c_{ij} = a_{ij} \cdot b_{ij}$ . PyTorch uses * or torch.mul(), with $O(mn)$ complexity—parallelizable and memory-bound.

Key Prerequisites:

Tensor shapes: MatMul needs (..., m, n) and (..., n, p); element-wise needs identical shapes.
Broadcasting: PyTorch auto-broadcasts for element-wise if compatible.
Integer inputs: No floating-point concerns here, but watch for overflow.

2. Algorithm Approach

Standard Approach: Use PyTorch primitives directly—no need to implement from scratch, as they are BLAS-optimized.

MatMul: torch.matmul(a, b) or a @ b. Handles broadcasting and batch dims.
Element-wise: torch.mul(a, b) or a * b.

Manual Fallback (Educational): For MatMul, nested loops over outer dims and dot product inner; for element-wise, loop over all elements. PyTorch vectorizes this via SIMD/AVX.

Output Conversion: Tests expect nested lists, not tensors: result.tolist().

3. Step-by-Step Strategy

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.