Introduction

In PyTorch, batch size is often dynamic. A training loader may emit batches of 64, while the final batch in an epoch might be smaller, and inference code might run on one example at a time. If you need to match another tensor to that batch dimension, hard-coding the repeat count will eventually break.

The solution is to read the batch dimension from the input tensor and tile or broadcast from there. The key design choice is whether you need real copies with repeat or a lightweight view with expand.

Reading the Batch Size at Runtime

Most of the time the batch dimension is x.size(0) or x.shape[0]. Once you have that value, you can reshape the source tensor so it has a singleton batch dimension and then replicate it.

For example, suppose you have a learnable template of shape (H, W) and want one copy per batch item:

1import torch
2
3x = torch.randn(5, 3, 32, 32)   # batch of 5 images
4template = torch.randn(4, 4)
5
6batch_size = x.size(0)
7tiled = template.unsqueeze(0).repeat(batch_size, 1, 1)
8
9print(tiled.shape)

The output shape is (5, 4, 4). Because batch_size comes from x, the code works for any batch length.

Prefer expand When Copies Are Not Needed

repeat creates actual repeated data in memory. That is fine when you need independent materialized values, but many batch-alignment tasks only need a broadcasted view.

In those cases, expand is usually better:

1import torch
2
3x = torch.randn(8, 16)
4bias = torch.arange(16, dtype=torch.float32)
5
6expanded = bias.unsqueeze(0).expand(x.size(0), -1)
7result = x + expanded
8
9print(expanded.shape)
10print(result.shape)

expand does not physically copy the tensor across the batch dimension. It reuses the same storage and pretends the singleton dimension has been stretched. That is memory-efficient and usually faster.

A Common Model Pattern

Imagine a module that compares each batch item with the same reference vector:

1import torch
2import torch.nn as nn
3
4class ReferenceMatcher(nn.Module):
5    def __init__(self, feature_dim):
6        super().__init__()
7        self.reference = nn.Parameter(torch.randn(feature_dim))
8
9    def forward(self, x):
10        ref = self.reference.unsqueeze(0).expand(x.size(0), -1)
11        return torch.cat([x, ref], dim=1)
12
13model = ReferenceMatcher(feature_dim=4)
14batch = torch.randn(3, 4)
15output = model(batch)
16print(output.shape)

The module does not care whether the incoming batch has 3, 32, or 128 rows. It just expands the reference parameter to match.

When repeat Is the Right Choice

Use repeat if later code needs actual copies that can diverge as separate tensors. One example is when you plan to reshape or mutate the repeated values independently in a way that cannot be expressed by broadcasting.

1import torch
2
3prototype = torch.tensor([[1.0, 2.0]])
4copies = prototype.repeat(4, 1)
5copies[0, 0] = 99.0
6
7print(prototype)
8print(copies)

Here the repeated tensor is separate materialized data. If you had used expand, an in-place write would either fail or behave in a way you did not intend.

tile, repeat, and Broadcasting

PyTorch also provides torch.tile, which can be convenient if you think in NumPy-style tiling patterns. Under the hood, though, the same conceptual question remains: do you want copied data or broadcasted behavior.

For many operations, the cleanest answer is to skip tiling entirely and let PyTorch broadcast automatically:

1import torch
2
3x = torch.randn(6, 10)
4scale = torch.linspace(1.0, 2.0, 10)
5
6y = x * scale
7print(y.shape)

Because scale has shape (10,), PyTorch automatically broadcasts it across the batch dimension. That is even simpler than calling expand.

Common Pitfalls

The most common mistake is using repeat everywhere. It works, but it can multiply memory usage quickly for large tensors. If the extra batch copies are read-only, expand or plain broadcasting is better.

Another pitfall is forgetting the singleton dimension. expand(batch_size, -1) only works after unsqueeze(0) creates a batch axis of size 1.

In-place writes are another source of bugs. Expanded tensors share storage, so they are not a good target for independent modification per batch item.

Finally, be careful about which dimension is the batch dimension. In most PyTorch models it is dimension 0, but sequence models and custom data layouts can differ. Always inspect the real input shape before tiling.

Summary

• Read the runtime batch size with x.size(0) or x.shape[0].
• Use unsqueeze before matching a per-sample tensor to the batch dimension.
• Prefer expand or plain broadcasting when you do not need real copies.
• Use repeat only when materialized repeated data is actually required.
• Dynamic tiling is easiest to maintain when it follows the actual input shape instead of hard-coded constants.