Introduction

tf.tensordot is a flexible tensor-contraction operator, but that flexibility is exactly why it causes trouble in batch multiplication code. The most common problems are contracting the wrong axes, collapsing the batch dimension by accident, and creating output shapes that are mathematically valid but operationally surprising.

If your real goal is ordinary batched matrix multiplication, tf.matmul is usually clearer. tensordot becomes useful when the contraction pattern is genuinely more general than matrix multiplication, but then you need to be very explicit about shapes and axes.

Understand What tensordot Actually Contracts

tf.tensordot(a, b, axes) sums over the selected axes of a and b. If those axes do not have equal lengths, TensorFlow throws a shape error.

1import tensorflow as tf
2
3a = tf.random.normal((32, 10, 64))
4b = tf.random.normal((32, 64, 128))
5
6out = tf.tensordot(a, b, axes=[[2], [1]])
7print(out.shape)

This is valid, but the output shape is not the usual batched matrix-multiplication shape. The batch dimension from both tensors remains in the result because you did not contract it.

That is the first mental trap: tensordot is not secretly "batched matmul with a different name." It is general contraction.

Use matmul for Standard Batched Matrix Multiplication

If you want the common batch operation batch x rows x cols multiplied by batch x cols x out, use tf.matmul:

1x = tf.random.normal((32, 10, 64))
2y = tf.random.normal((32, 64, 128))
3z = tf.matmul(x, y)
4
5print(z.shape)  # (32, 10, 128)

This keeps the batch dimension intact and does exactly what most people mean by batch multiplication.

If you force the same idea through tensordot, you can easily get a higher-rank tensor that needs reshaping or whose semantics are unclear.

Avoid Contracting the Batch Axis by Accident

One of the most common tensordot bugs is accidentally reducing over axis 0, which is often the batch dimension:

# Usually wrong for batch-independent multiplication
bad = tf.tensordot(a, b, axes=[[0, 2], [0, 1]])

This contracts the batch axis along with the feature axis. The result is no longer "one output per batch item." It is a cross-batch contraction.

If each batch element should remain independent, do not reduce the batch axis.

Use einsum When the Shape Story Needs to Be Readable

tf.einsum is often easier to reason about because the preserved and contracted axes are visible in the equation:

result = tf.einsum("bij,bjk->bik", a, b)
print(result.shape)  # (32, 10, 128)

For many advanced contractions, einsum is easier to debug than a nested axes=[[...], [...]] specification because the equation makes the intent more obvious.

Add Shape Assertions Before Expensive Ops

Shape bugs are cheaper to catch early:

tf.debugging.assert_equal(tf.shape(a)[0], tf.shape(b)[0])  # same batch size
tf.debugging.assert_equal(tf.shape(a)[2], tf.shape(b)[1])  # inner dimensions

These checks are especially useful in custom layers, where a later refactor may change the tensor layout without updating the contraction logic.

Watch Memory and Kernel Choice

General tensor contractions can allocate large intermediate buffers. That means a mathematically correct tensordot may still be a bad operational choice for large inputs.

A few practical guidelines:

• prefer matmul for matrix-like workloads because it maps to optimized kernels
• consider einsum when readability helps avoid axis mistakes
• profile memory usage for high-rank contractions
• chunk very large problems if the contraction explodes memory

For performance-sensitive code, measurement matters more than API preference.

Common Pitfalls

The biggest mistake is using tensordot for ordinary batched matrix multiplication when matmul expresses the intent more directly.

Another common issue is reducing the batch axis unintentionally. This changes the semantics of the operation, not just the output shape.

People also assume the result shape will "look right" automatically. With general contraction, that assumption is unsafe.

Finally, do not wait until the model fails deep in training to inspect tensor shapes. Add assertions and small debug prints near the contraction itself.

Summary

• 'tf.tensordot performs general tensor contraction, not automatic batched matrix multiplication.'
• Use tf.matmul for standard batched matrix products.
• Keep the batch axis out of the contracted axes unless you truly want cross-batch reduction.
• Use tf.einsum when you need the contraction logic to be more readable.
• Add shape checks and profile memory for high-rank contractions.