TensorFlow
covariance matrix
machine learning
data science
Python

how to get covariance matrix in tensorflow?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

A covariance matrix summarizes how features vary together across a dataset. TensorFlow does not require a special high-level covariance API to compute it. The standard approach is to center the data by subtracting the mean and then multiply the centered matrix by its transpose with the correct normalization factor.

Decide Which Axis Represents Features

Before writing code, decide how your data is shaped. A common convention is:

  • rows are samples
  • columns are features

If X has shape n x d, then the covariance matrix has shape d x d because it measures relationships between features.

That convention is what the examples below assume.

Compute Covariance from First Principles

The sample covariance matrix is:

  • center each feature by subtracting its mean
  • multiply the centered matrix transpose by the centered matrix
  • divide by n - 1

In TensorFlow:

python
1import tensorflow as tf
2
3X = tf.constant([
4    [1.0, 2.0],
5    [2.0, 1.0],
6    [3.0, 4.0],
7    [4.0, 3.0],
8], dtype=tf.float32)
9
10mean = tf.reduce_mean(X, axis=0, keepdims=True)
11X_centered = X - mean
12n = tf.cast(tf.shape(X)[0], tf.float32)
13
14cov = tf.matmul(X_centered, X_centered, transpose_a=True) / (n - 1.0)
15print(cov.numpy())

This returns a 2 x 2 covariance matrix because the input has two features.

Sample Covariance vs Population Covariance

The normalization factor depends on what you want.

  • sample covariance divides by n - 1
  • population covariance divides by n

If you want the population version, change the final line:

python
cov_population = tf.matmul(X_centered, X_centered, transpose_a=True) / n

This is a statistics decision, not a TensorFlow-specific one. The code structure is the same either way.

Wrap It in a Reusable Function

A helper makes the intent clearer and avoids repeating the same math.

python
1import tensorflow as tf
2
3
4def covariance_matrix(X, sample=True):
5    X = tf.convert_to_tensor(X, dtype=tf.float32)
6    mean = tf.reduce_mean(X, axis=0, keepdims=True)
7    X_centered = X - mean
8    n = tf.cast(tf.shape(X)[0], tf.float32)
9    denom = n - 1.0 if sample else n
10    return tf.matmul(X_centered, X_centered, transpose_a=True) / denom
11
12
13X = [[1.0, 2.0], [2.0, 1.0], [3.0, 4.0], [4.0, 3.0]]
14print(covariance_matrix(X).numpy())

This is usually enough for machine-learning preprocessing, exploratory analysis, or covariance-based regularization experiments.

Covariance Between Two Variables Only

If you only need the covariance between two one-dimensional tensors, you can compute it directly.

python
1import tensorflow as tf
2
3x = tf.constant([1.0, 2.0, 3.0, 4.0])
4y = tf.constant([2.0, 1.0, 4.0, 3.0])
5
6x_centered = x - tf.reduce_mean(x)
7y_centered = y - tf.reduce_mean(y)
8
9cov_xy = tf.reduce_sum(x_centered * y_centered) / (tf.cast(tf.shape(x)[0], tf.float32) - 1.0)
10print(cov_xy.numpy())

That value would appear in an off-diagonal entry of the full covariance matrix.

Batches and Large Tensors

For large datasets, the main consideration is memory, not the formula itself. If the full n x d tensor is too large, you may need a streamed or batched covariance computation instead of materializing the entire matrix at once.

For moderate tensor sizes, TensorFlow handles the matrix operations efficiently on CPU or GPU. The core formula remains the same.

Correlation Is Different

A covariance matrix is not the same as a correlation matrix. Correlation rescales each feature by its standard deviation so the values are normalized between -1 and 1 in the usual interpretation.

If your downstream task expects correlation, do not stop at covariance.

Common Pitfalls

The most common mistake is mixing up which axis represents samples and which axis represents features, which leads to a matrix of the wrong shape. Another is forgetting to center the data before multiplying, which produces a second-moment matrix rather than covariance. Developers also often divide by n when they intended sample covariance with n - 1. A final issue is assuming covariance and correlation are interchangeable even though they answer different statistical questions.

Summary

  • In TensorFlow, covariance is usually computed by centering the data and using tf.matmul.
  • With n x d input, the covariance matrix is typically d x d.
  • Use n - 1 for sample covariance and n for population covariance.
  • Always subtract the feature means before computing covariance.
  • Be explicit about axis conventions so the result matches the analysis you intend.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.