TensorFlow
Tensor Concatenation
Machine Learning
Deep Learning
Data Preprocessing

How to concatenate two tensors horizontally 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

Concatenating tensors is a fundamental operation in TensorFlow, used frequently when merging feature vectors, combining model outputs, or preparing input batches. Horizontal concatenation joins tensors along columns (axis 1 for 2D tensors), while vertical concatenation joins along rows (axis 0). Understanding how tf.concat works with different axes and tensor shapes prevents common runtime errors.

Horizontal Concatenation with tf.concat

The primary function for combining tensors is tf.concat. To concatenate two 2D tensors side by side (horizontally), set axis=1.

python
1import tensorflow as tf
2
3a = tf.constant([[1, 2], [3, 4]])      # shape (2, 2)
4b = tf.constant([[5, 6], [7, 8]])      # shape (2, 2)
5
6c = tf.concat([a, b], axis=1)          # shape (2, 4)
7print(c)
text
tf.Tensor(
[[1 2 5 6]
 [3 4 7 8]], shape=(2, 4), dtype=int32)

The result has the same number of rows as the inputs, but the columns are combined. Both tensors must have the same number of rows for this to work.

Vertical Concatenation with axis=0

Setting axis=0 stacks tensors vertically, combining rows.

python
d = tf.concat([a, b], axis=0)          # shape (4, 2)
print(d)
text
1tf.Tensor(
2[[1 2]
3 [3 4]
4 [5 6]
5 [7 8]], shape=(4, 2), dtype=int32)

For vertical concatenation, the column count must match across all input tensors.

Concatenating Along Higher Dimensions

With 3D tensors (common in batch processing and sequence models), the axis parameter selects which dimension to join.

python
1# Two batches of 2x3 matrices
2x = tf.constant([[[1, 2, 3], [4, 5, 6]]])   # shape (1, 2, 3)
3y = tf.constant([[[7, 8, 9], [10, 11, 12]]]) # shape (1, 2, 3)
4
5# Concatenate along batch dimension
6batch_concat = tf.concat([x, y], axis=0)     # shape (2, 2, 3)
7
8# Concatenate along row dimension
9row_concat = tf.concat([x, y], axis=1)       # shape (1, 4, 3)
10
11# Concatenate along column dimension
12col_concat = tf.concat([x, y], axis=2)       # shape (1, 2, 6)

Using negative indexing also works. axis=-1 always refers to the last dimension, which is equivalent to horizontal concatenation for 2D tensors.

python
e = tf.concat([a, b], axis=-1)  # same as axis=1 for 2D
print(e.shape)                  # (2, 4)

tf.concat vs tf.stack

While tf.concat joins tensors along an existing axis, tf.stack creates a new axis and stacks the tensors along it. This changes the output rank.

python
1# tf.concat: no new dimension
2concat_result = tf.concat([a, b], axis=0)
3print(concat_result.shape)   # (4, 2)
4
5# tf.stack: adds a new dimension
6stack_result = tf.stack([a, b], axis=0)
7print(stack_result.shape)    # (2, 2, 2)

Use tf.stack when you want to create a batch dimension from individual samples. Use tf.concat when you want to extend an existing dimension.

python
1# Creating a batch from individual images
2img1 = tf.random.normal([224, 224, 3])
3img2 = tf.random.normal([224, 224, 3])
4
5batch = tf.stack([img1, img2], axis=0)
6print(batch.shape)  # (2, 224, 224, 3)

Handling Shape Mismatches

TensorFlow raises an InvalidArgumentError when dimensions other than the concatenation axis do not match. You can pad or reshape tensors before concatenating.

python
1# Tensors with different column counts
2p = tf.constant([[1, 2]])       # shape (1, 2)
3q = tf.constant([[3, 4, 5]])    # shape (1, 3)
4
5# This fails: tf.concat([p, q], axis=0) because columns differ
6# Fix by padding the shorter tensor
7p_padded = tf.pad(p, [[0, 0], [0, 1]])  # pad 1 column on the right
8print(p_padded)                          # [[1, 2, 0]]
9
10result = tf.concat([p_padded, q], axis=0)  # shape (2, 3)
11print(result)
text
tf.Tensor(
[[1 2 0]
 [3 4 5]], shape=(2, 3), dtype=int32)

Concatenating Ragged Tensors

When working with sequences of variable length, tf.ragged provides specialized support.

python
1ragged_a = tf.ragged.constant([[1, 2], [3, 4, 5]])
2ragged_b = tf.ragged.constant([[6], [7, 8]])
3
4ragged_concat = tf.concat([ragged_a, ragged_b], axis=0)
5print(ragged_concat)
text
<tf.RaggedTensor [[1, 2], [3, 4, 5], [6], [7, 8]]>

Ragged tensors allow concatenation without requiring uniform shapes on non-concat dimensions, which is useful for NLP and variable-length sequence processing.

Common Pitfalls

  • Confusing axis=0 and axis=1 leads to unexpected output shapes; axis=1 is horizontal (column-wise) for 2D tensors, while axis=0 is vertical (row-wise).
  • Mixing tf.concat with tf.stack without realizing that tf.stack adds a new dimension while tf.concat extends an existing one, resulting in unexpected tensor ranks.
  • Rank mismatch between inputs (for example, concatenating a 2D tensor with a 3D tensor) causes an immediate error; reshape or expand dimensions first with tf.expand_dims.
  • Forgetting that all non-concat dimensions must match exactly is the most common shape error; use tf.pad or tf.reshape to align shapes before concatenation.
  • Not specifying the axis parameter defaults to axis=0, which is vertical concatenation; always pass the axis explicitly for clarity.

Summary

  • Use tf.concat([a, b], axis=1) for horizontal concatenation and axis=0 for vertical concatenation of 2D tensors.
  • Use axis=-1 as a portable shorthand for concatenating along the last dimension regardless of tensor rank.
  • Choose tf.stack over tf.concat when you need to create a new dimension, such as forming a batch from individual samples.
  • All dimensions except the concatenation axis must match; use tf.pad or tf.reshape to fix shape mismatches before concatenation.
  • For variable-length data, use tf.ragged tensors, which support concatenation without requiring uniform shapes on non-concat dimensions.

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.