Neural-network parameters are usually matrices (weight matrices $W$). To train them we need the gradient of a loss with respect to a matrix. Differentiating a matrix by a vector — or a matrix by a matrix — produces a higher-dimensional object: a tensor. This page is about keeping the bookkeeping straight, because the shapes are where everyone gets confused.

A real-life example: nudging a weight matrix

A dense layer computes $\mathbf{y} = W\mathbf{x}$. During training we ask: “if I nudge each entry $W_{ij}$ a little, how does the loss change?” The answer is a gradient with one number per entry of $W$ — the same shape as $W$ itself. That’s what an optimizer adds (scaled) back to $W$. But deriving it carefully means tracking how a whole matrix output responds to a whole matrix input.

The shape of a derivative

The most important skill here is predicting the shape of a gradient before computing it. For a function $f : \mathbb{R}^D \to \mathbb{R}^E$, the gradient (Jacobian) is $E \times D$:

See how the shape depends on dimensions

The gradient’s shape is determined entirely by the input and output dimensions. Watch the gradient grid morph as the function’s input size $D$ and output size $E$ change — scalar → row → column → full matrix:

Differentiating by a matrix gives a tensor

When you differentiate an $m \times n$ matrix $A$ with respect to a vector $\mathbf{x} \in \mathbb{R}^p$, you get a 3rd-order tensor of shape $(m \times n) \times p$ — one matrix of partials per input component. Differentiating a matrix by a matrix gives a 4th-order tensor. These are just multi-dimensional arrays; the entry $J_{ijkl} = \partial A_{ij} / \partial B_{kl}$.

The flattening trick

Tensors are awkward. A common shortcut exploits the isomorphism $\mathbb{R}^{m\times n} \cong \mathbb{R}^{mn}$: flatten the matrix into a long vector (stack its columns), compute an ordinary Jacobian matrix, then reshape the result back into a tensor. Flattening turns the chain rule back into plain matrix multiplication — which is why deep-learning frameworks store parameters as flattened vectors under the hood.

  flowchart TD
  M["gradient of matrix A wrt vector x"] --> T["3rd-order tensor (m×n)×p"]
  T --> F["flatten A → vector (mn)"]
  F --> J["ordinary Jacobian (mn)×p"]
  J --> R["reshape back to tensor"]
  R -.-> BP["frameworks flatten params → matrix chain rule"]

NumPy: shapes and the flatten trick

import numpy as np
 
# f(x) = A x, with A (M×N), x (N,). Gradient df/dx is M×N = A.
M, N = 3, 4
A = np.arange(M*N).reshape(M, N).astype(float)
x = np.ones(N)
 
def jac(f, x, h=1e-6):
    fx = f(x); J = np.zeros((len(fx), len(x)))
    for j in range(len(x)):
        xp = x.copy(); xp[j] += h
        J[:, j] = (f(xp) - fx) / h
    return J
 
print("df/dx shape:", jac(lambda v: A @ v, x).shape, "== A shape:", A.shape)
 
# Gradient of a matrix output wrt a vector input -> a 3rd-order tensor.
# g(t) = t * A  (scale A by scalar t). dG/dt has the same shape as A.
def G(t): return t * A
dGdt = (G(1.0 + 1e-6) - G(1.0)) / 1e-6
print("dG/dt shape:", dGdt.shape, "(a matrix, = ∂ each entry / ∂t)")
 
# Flatten trick: reshape A (3×4) to a length-12 vector and back
flat = A.reshape(-1)
print("flattened length:", flat.shape[0], "-> reshaped back equal:", np.allclose(flat.reshape(M, N), A))

import numpy as np
 
# f(x) = A x, with A (M×N), x (N,). Gradient df/dx is M×N = A.
M, N = 3, 4
A = np.arange(M*N).reshape(M, N).astype(float)
x = np.ones(N)
 
def jac(f, x, h=1e-6):
    fx = f(x); J = np.zeros((len(fx), len(x)))
    for j in range(len(x)):
        xp = x.copy(); xp[j] += h
        J[:, j] = (f(xp) - fx) / h
    return J
 
print("df/dx shape:", jac(lambda v: A @ v, x).shape, "== A shape:", A.shape)
 
# Gradient of a matrix output wrt a vector input -> a 3rd-order tensor.
# g(t) = t * A  (scale A by scalar t). dG/dt has the same shape as A.
def G(t): return t * A
dGdt = (G(1.0 + 1e-6) - G(1.0)) / 1e-6
print("dG/dt shape:", dGdt.shape, "(a matrix, = ∂ each entry / ∂t)")
 
# Flatten trick: reshape A (3×4) to a length-12 vector and back
flat = A.reshape(-1)
print("flattened length:", flat.shape[0], "-> reshaped back equal:", np.allclose(flat.reshape(M, N), A))

df/dx shape: (3, 4) == A shape: (3, 4)
dG/dt shape: (3, 4) (a matrix, = ∂ each entry / ∂t)
flattened length: 12 -> reshaped back equal: True

df/dx shape: (3, 4) == A shape: (3, 4)
dG/dt shape: (3, 4) (a matrix, = ∂ each entry / ∂t)
flattened length: 12 -> reshaped back equal: True

Why this matters for ML

• Weight-matrix gradients are what optimizers update; frameworks flatten them so updates are simple vector arithmetic.
• Understanding shapes prevents the most common deep-learning bug: dimension mismatches in hand-derived or custom gradients.
• Tensor derivatives appear whenever a layer’s output and parameters are both matrices (attention, convolutions).

🧪 Try It Yourself

Exercise 1 – Predict the gradient shape

Exercise 2 – Gradient of A x is A

Exercise 3 – Flatten and reshape

Recap

• The gradient of $f : \mathbb{R}^D \to \mathbb{R}^E$ has shape $E \times D$ — always predict the shape first.
• Differentiating a matrix by a vector/matrix produces a tensor (3rd/4th order).
• The flatten trick ($\mathbb{R}^{m\times n} \cong \mathbb{R}^{mn}$) turns tensor derivatives back into ordinary Jacobian matrices — how frameworks handle weight gradients.

Next: a cheat sheet of ready-made matrix-gradient formulas — Useful Identities.