Easy Framework

Concatenate Tensors

Concatenation joins tensors along an existing axis. Here, both inputs and the result are matrices: no new axis is added.

Your task

Implement concat(a, b, axis) for two-dimensional tensors with the same dtype and device. The supplied axis is either 0 or 1, and the other dimension matches. Return the joined tensor, preserving dtype and device. Leave both inputs unchanged.

With axis=0, append b's rows after a's rows. Shapes (m, n) and (k, n) produce (m + k, n).
With axis=1, append b's columns after a's columns in each row. Shapes (m, n) and (m, k) produce (m, n + k).

For a=[[1, 2]] and b=[[3, 4]], axis=0 gives [[1, 2], [3, 4]], while axis=1 gives [[1, 2, 3, 4]].

The inputs need not have identical shapes. For example, joining [[1], [2]] and [[3, 4], [5, 6]] along axis 1 gives [[1, 3, 4], [2, 5, 6]]: only the row counts need to match.

Use torch.cat([a, b], dim=axis) or jnp.concatenate([a, b], axis=axis). Keep a before b along the chosen axis. Do not flatten the inputs or use stack, which would introduce a new dimension.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Give each tensor axis a meaning

A tensor's shape tells you how many entries each axis has. It does not tell you what those entries represent. For example, (batch, features) and (features, batch) use different axis orders. Name the axes before choosing an operation.

A reduction combines entries along an axis, for example by summing or averaging. Averaging the features of a (B, D) tensor gives one number per example, with shape (B,). Keeping the reduced axis gives shape (B, 1) instead.

For [[2, 4, 6], [1, 3, 5]], the row means with the axis kept are [[4], [3]]. Subtracting them from the input gives [[-2, 0, 2], [-2, 0, 2]]: each row has been centered using its own mean.

Reshape and transpose do different jobs

A reshape regroups elements using the operation's traversal order. A transpose changes the axis order. For the usual row-major reshape, [[1, 2, 3], [4, 5, 6]] reshaped to (3, 2) becomes [[1, 2], [3, 4], [5, 6]]. Transposing the original instead gives [[1, 4], [2, 5], [3, 6]].

Both results have shape (3, 2), but their entries mean different things. Check values as well as shapes. Unequal axis sizes make many axis mistakes easier to spot than square examples do.

Concatenate along an axis or stack on a new one

Concatenation extends an existing axis. Stacking creates a new axis. For two arrays of shape (2, 3), concatenating on axis zero gives (4, 3). Stacking on a new leading axis gives (2, 2, 3).

Detecting runtime…

Concatenate Tensors

Your task

Give each tensor axis a meaning

Reshape and transpose do different jobs

Concatenate along an axis or stack on a new one

Your code, saved here

Stuck?

JAX reference solution