Introduction

tf.matmul, tf.tensordot, and tf.einsum all express tensor contraction, but they are not interchangeable from a readability point of view. The right choice depends on how standard the operation is, how many axes are involved, and whether the formula is easier to understand as axis lists or as equation notation.

Use matmul for Ordinary Matrix Products

If you are doing a normal matrix multiply, or a batched version of one, tf.matmul should usually be the default. It communicates intent immediately.

1import tensorflow as tf
2
3a = tf.constant([[1., 2.], [3., 4.]])
4b = tf.constant([[5., 6.], [7., 8.]])
5
6c = tf.matmul(a, b)
7print(c.numpy())

It also supports transpose flags, which covers many common model-building cases without reshaping tensors by hand.

1x = tf.random.normal((4, 8, 16))
2y = tf.random.normal((4, 16, 32))
3
4batched = tf.matmul(x, y)
5print(batched.shape)  # (4, 8, 32)

For layers, projections, and batched matrix math, matmul is usually the clearest API.

Use tensordot When You Want to Name the Contracted Axes

tf.tensordot is useful when the operation is not a classic matrix multiply but still has a straightforward axis-contraction story.

1x = tf.random.normal((2, 3, 4))
2y = tf.random.normal((4, 5))
3
4z = tf.tensordot(x, y, axes=[[2], [0]])
5print(z.shape)

Here the last axis of x is contracted with the first axis of y, producing shape (2, 3, 5). tensordot is especially readable when you want to say "sum over these axes" and nothing more.

You can also contract multiple axes at once:

1a = tf.random.normal((2, 3, 4))
2b = tf.random.normal((3, 4, 6))
3
4out = tf.tensordot(a, b, axes=[[1, 2], [0, 1]])
5print(out.shape)  # (2, 6)

Use einsum for Formulas with Several Axis Relationships

tf.einsum is the most expressive option. Instead of passing axis lists, you describe the relationship between dimensions with symbols.

1m = tf.random.normal((2, 3))
2n = tf.random.normal((3, 4))
3
4p = tf.einsum('ik,kj->ij', m, n)
5print(p.shape)

That is mathematically compact and extremely powerful. It becomes especially useful for attention-style code and cases where multiplication, summation, and axis reordering all happen together.

1q = tf.random.normal((2, 4, 8))
2k = tf.random.normal((2, 6, 8))
3
4scores = tf.einsum('bqd,bkd->bqk', q, k)
5print(scores.shape)  # (2, 4, 6)

In that example, the equation says more clearly what is being compared than a sequence of reshapes and transposes would.

Compare Equivalent Forms

Many operations can be written with more than one API.

1a = tf.random.normal((5, 7))
2b = tf.random.normal((7, 9))
3
4r1 = tf.matmul(a, b)
5r2 = tf.tensordot(a, b, axes=[[1], [0]])
6r3 = tf.einsum('ab,bc->ac', a, b)
7
8print(tf.reduce_max(tf.abs(r1 - r2)).numpy())
9print(tf.reduce_max(tf.abs(r1 - r3)).numpy())

All three operations produce the same result here. That does not mean they are equally readable. For a plain two-dimensional product, matmul is the strongest choice because it is obvious at a glance.

A Practical Decision Rule

Use matmul when the operation is fundamentally matrix multiplication.

Use tensordot when you need to contract specific axes and the mapping is still simple enough to explain with lists.

Use einsum when the expression has several axis relationships and equation notation is easier to read than a mix of transpose, reshape, and multiply calls.

That choice is mostly about maintainability, not just performance. In practice, shape bugs cost more time than tiny differences in API overhead.

Debug Shapes Before Optimizing

Most failures come from mismatched dimensions or incorrect assumptions about batch axes. Print shapes before the operation and verify the algebra on paper if necessary.

print("a:", a.shape)
print("b:", b.shape)

For einsum, check that each symbol means exactly one dimension and that repeated symbols are the axes you intend to sum over.

Profile Only with Realistic Tensor Shapes

Performance depends on tensor shapes, backend kernels, device type, and graph optimizations. Benchmark only if performance matters for your actual workload.

1import time
2
3left = tf.random.normal((128, 512))
4right = tf.random.normal((512, 512))
5
6start = time.time()
7for _ in range(200):
8    _ = tf.matmul(left, right)
9print("matmul seconds:", time.time() - start)

Do not assume einsum is always slower or always faster. TensorFlow may lower different forms to similar kernels.

Common Pitfalls

• Using einsum for a simple two-dimensional multiply and making the code harder to read.
• Passing the wrong axis lists to tensordot and silently producing an unexpected shape.
• Forgetting that matmul treats leading dimensions as batch dimensions.
• Comparing performance on toy tensors and then generalizing the result to production workloads.
• Reading an einsum equation incorrectly because the symbolic dimensions are not documented.

Summary

• Prefer tf.matmul for standard matrix and batched matrix multiplication.
• Reach for tf.tensordot when the job is explicit axis contraction.
• Use tf.einsum when equation notation makes a multi-axis formula clearer.
• Validate shapes before assuming an API call is wrong.
• Benchmark only with realistic shapes and hardware if performance is important.