You can use torch.expand

t = torch.ones((1, 1000, 1000))
t10 = t.expand(10, 1000, 1000)

Keep in mind that the t10 is just a reference to t. So for example, a change to t10[0,0,0] will result in the same change in t[0,0,0] and every member of t10[:,0,0].

Other than direct access, most operations performed on t10 will cause memory to be copied which will break the reference and cause more memory to be used. For example: changing the device (.cpu(), .to(device=...), .cuda()), changing the datatype (.float(), .long(), .to(dtype=...)), or using .contiguous().

tensor.repeat should suit your needs but you need to insert a unitary dimension first. For this we could use either tensor.unsqueeze or tensor.reshape. Since unsqueeze is specifically defined to insert a unitary dimension we will use that.

B = A.unsqueeze(1).repeat(1, K, 1)

Code Description A.unsqueeze(1) turns A from an [M, N] to [M, 1, N] and .repeat(1, K, 1) repeats the tensor K times along the second dimension.

Use this:

import torch
# A is your tensor
B = torch.tensor([1, 2, 3])
C = A.repeat_interleave(B, dim = 0)

EDIT:

The above works fine if A is a single 2D tensor. To repeat all (2D) tensors in a batch in the same manner, this is a simple workaround:

A = torch.tensor([[[1, 1, 1], [2, 2, 2], [3, 3, 3]], 
    [[1, 2, 3], [4, 5, 6], [2,2,2]]]) # A has 2 tensors each of shape (3, 3)
B = torch.tensor([1, 2, 3]) # Rep. of each row of every tensor in the batch

A1 = A.reshape(1, -1, A.shape[2]).squeeze()
B1 = B.repeat(A.shape[0])
C = A1.repeat_interleave(B1, dim = 0).reshape(A.shape[0], -1, A.shape[2])

C is:

tensor([[[1, 1, 1],
         [2, 2, 2],
         [2, 2, 2],
         [3, 3, 3],
         [3, 3, 3],
         [3, 3, 3]],

        [[1, 2, 3],
         [4, 5, 6],
         [4, 5, 6],
         [2, 2, 2],
         [2, 2, 2],
         [2, 2, 2]]])

As you can see each inside tensor in the batch is repeated in the same manner.

Assuming that the first dimension of a is 1 as in your example, you could do the following:

a = torch.Tensor([[1, 2, 3, 4]])
b_abbreviated = torch.Tensor([[1, 2, 3], [4, 5, 6]])
torch.mm(a.reshape(-1, 2), b_abbreviated).sum(axis=0, keepdim=True)

Here, instead of repeating the rows, you multiply a in chunks, then add them up column-wise to get the same result.

If the first dimension of a is not necessarily 1, you could try the following:

torch.cat(torch.split(torch.mm(a.reshape(-1,2),b_abbreviated), a.shape[0]), dim=1).sum(
dim=0, keepdim=True).reshape(a.shape[0], -1)

Here, you do the following:

• With torch.mm(a.reshape(-1,2),b_abbreviated, you again split each row of a into chunks of size 2 and stack them one over the other, and then stack each row over the other.
• With torch.split(torch.mm(a.reshape(-1,2),b_abbreviated), a.shape[0]), these stacks are then separated row-wise, so that each resultant component of the split corresponds to chunks of a single row.
• With torch.cat(torch.split(torch.mm(a.reshape(-1,2),b_abbreviated), a.shape[0]), dim=1) these stacks are then concatenated column-wise.
• With .sum(dim=0, keepdim=True), results corresponding to separate chunks of individual rows in a are added up.
• With .reshape(a.shape[0], -1), rows of a that were concatenated column-wise are again stacked row-wise.

It seems quite slow compared to direct matrix multiplication, which is not surprising, but I have not yet checked in comparison to explicit iteration. There are likely better ways of doing this, will edit if I think of any.

Try torch.repeat_interleave() method: https://pytorch.org/docs/stable/torch.html#torch.repeat_interleave

>>> x = torch.tensor([1, 2, 3])
>>> x.repeat_interleave(2)
tensor([1, 1, 2, 2, 3, 3])