Introduction

In TensorFlow, a rotation matrix is usually built from tf.cos and tf.sin so the result stays differentiable and can be used inside training pipelines. The main concerns are getting the matrix formula right, handling batch dimensions correctly, and applying the transform without shape errors.

The exact matrix depends on whether you are rotating in 2D or 3D. For most machine learning and computer vision tasks, a 2D rotation matrix is the simplest place to start.

A 2D Rotation Matrix

A counterclockwise rotation by angle theta uses:

• first row: cos(theta), -sin(theta)
• second row: sin(theta), cos(theta)

In TensorFlow:

1import tensorflow as tf
2
3
4def rotation_matrix_2d(theta):
5    c = tf.cos(theta)
6    s = tf.sin(theta)
7    return tf.stack([
8        tf.stack([c, -s], axis=-1),
9        tf.stack([s,  c], axis=-1),
10    ], axis=-2)
11
12
13theta = tf.constant(0.5, dtype=tf.float32)
14R = rotation_matrix_2d(theta)
15print(R)

This produces a 2 x 2 tensor that can rotate points in the plane.

For a single angle, the returned shape is just 2 x 2. For a batch of angles, the same construction generalizes naturally to [batch, 2, 2], which is why building the matrix from TensorFlow ops instead of Python lists is so useful.

Applying the Matrix to a Point

point = tf.constant([1.0, 0.0], dtype=tf.float32)
rotated = tf.linalg.matvec(R, point)
print(rotated)

tf.linalg.matvec is a clean way to apply a matrix to a vector without manual reshaping.

Batched Angles

In ML workflows, you often need one rotation per sample. Instead of looping in Python, build all matrices at once:

1angles = tf.constant([0.0, 0.5, 1.0], dtype=tf.float32)
2
3c = tf.cos(angles)
4s = tf.sin(angles)
5
6R_batch = tf.stack([
7    tf.stack([c, -s], axis=-1),
8    tf.stack([s,  c], axis=-1),
9], axis=-2)
10
11points = tf.constant([
12    [1.0, 0.0],
13    [1.0, 0.0],
14    [1.0, 0.0],
15], dtype=tf.float32)
16
17rotated = tf.einsum("bij,bj->bi", R_batch, points)
18print(rotated)

This keeps the computation vectorized and accelerator-friendly.

A 3D Example Around the Z Axis

For 3D, the matrix depends on the chosen axis. A rotation around the z-axis looks like:

1def rotation_matrix_z(theta):
2    c = tf.cos(theta)
3    s = tf.sin(theta)
4    zero = tf.zeros_like(theta)
5    one = tf.ones_like(theta)
6
7    return tf.stack([
8        tf.stack([c, -s, zero], axis=-1),
9        tf.stack([s,  c, zero], axis=-1),
10        tf.stack([zero, zero, one], axis=-1),
11    ], axis=-2)
12
13
14theta = tf.constant(0.3, dtype=tf.float32)
15Rz = rotation_matrix_z(theta)
16print(Rz)

Once you have axis-specific matrices, you can compose them with tf.matmul.

Why This Works Well in TensorFlow

Because the matrix is built from TensorFlow ops, gradients flow through it automatically. That matters when the angle is a trainable variable.

1theta = tf.Variable(0.1, dtype=tf.float32)
2point = tf.constant([1.0, 0.0], dtype=tf.float32)
3target = tf.constant([0.0, 1.0], dtype=tf.float32)
4
5optimizer = tf.keras.optimizers.Adam(0.05)
6
7for _ in range(100):
8    with tf.GradientTape() as tape:
9        pred = tf.linalg.matvec(rotation_matrix_2d(theta), point)
10        loss = tf.reduce_sum((pred - target) ** 2)
11    grad = tape.gradient(loss, [theta])
12    optimizer.apply_gradients(zip(grad, [theta]))

That is why it is better to build the matrix with TensorFlow primitives instead of with NumPy inside a training graph.

Common Pitfalls

• Mixing degrees and radians. TensorFlow trig functions expect radians.
• Putting the sine signs in the wrong positions and silently rotating the wrong direction.
• Ignoring batch dimensions and getting matrix-shape mismatches.
• Using NumPy matrix construction in the middle of a differentiable TensorFlow pipeline.
• Forgetting that rotation order matters when composing multiple 3D rotations.

Summary

• A TensorFlow rotation matrix is built from tf.cos and tf.sin.
• Use a 2 x 2 matrix for 2D and axis-specific 3 x 3 matrices for 3D.
• 'tf.linalg.matvec and tf.einsum are practical ways to apply rotations.'
• Vectorize batched rotations instead of looping in Python.
• Build matrices with TensorFlow ops so gradients remain available.