To assign a tensor to a part of another tensor

Assume I have a 3*3 tensor a which is

a = torch.rand(3,3)
tensor([[0.5456, 0.2233, 0.1322],
        [0.6037, 0.4913, 0.6274],
        [0.6362, 0.7671, 0.6741]])

and a 2*2 tensor b which is

b = torch.tensor([[10,20], [30,40]])
tensor([[10, 20],
        [30, 40]])

How can I assign tensor b onto specific position on tensor a, for example I want to get like this efficiently (without for loop)

tensor([[10.0000,  0.2233, 20.0000],
        [ 0.6037,  0.4913,  0.6274],
        [30.0000,  0.7671, 40.0000]])

I’ve tried using a mask to extract a part of a, and assign b to a view of tensor a like below, but seems failed.

mask: tensor([[1, 0, 1],
              [0, 0, 0],
              [1, 0, 1]], dtype=torch.uint8)
a[mask] = b
RuntimeError: expand(torch.LongTensor{[2, 2]}, size=[4]): the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)

or

c = a[mask].view(2,2)
c = b

but seems c and a do not share memory

Thanks!:grinning:


Supplement: for each iteration, I always want to put tensor b into the same position.
In this example, I always want to put a 2 * 2 tensor into the four corners of a 3 * 3 tensor

Hi,

How do you know where you want to put them originaly? Do you have a mask? if so does it always have the same number of ones in each row and each column? Or indices where you want to put these values?

Thanks!
For each iteration, I always want to put tensor b into the same position.
In this example, I always want to put a 2 * 2 tensor into the four corners of a 3 * 3 tensor

Hi,

If you want to change the corners of a square matrix with 4 values, the most efficient way is to change the stride by hand to get the tensor you want:

import torch

# a must be 2x2
# b can be any squared matrix
# You can generalize to batch by keeping original size and stride below
a = torch.rand(2, 2)
b = torch.zeros(3, 3)
print("a")
print(a)
print("b")
print(b)


assert b.ndimension() == 2
assert b.is_contiguous()
size0, size1 = b.size()
stride0, stride1 = b.stride()
new_size = (2, 2)
new_stride = (stride0*(size0-1), stride1*(size1-1))

b_corner = b.as_strided(new_size, new_stride)
# b_corner is a "real" tensor that shares data with b

b_corner += a
print("b_corner")
print(b_corner)

print("b")
print(b)

More generally the approach would be:

  • If you can use it with custom stride (elements are always spaced by the same distance in a given dimension), do it this way.
  • If it’s really sparse use indices and put_ to set the values at the given indices. If you have more than one dimension to index, you might want to use scatter_.
  • If it’s fairly dense, then a masked approach using masked_scatter_ is the way to go (similar to what you tried). The point is that the source should be a 1D Tensor containing as many values as 1s in the mask. In your case a[mask] = b.view(-1) for example but you will need to be careful in which order you put the elements in b so that they are placed where you want to in the original matrix.